Out of memory issues using generators

Hi everyone,

I’m currently facing a memory issue in PyTorch, and couldn’t find an appropriate answer from existing topics (feel free to correct me on this one!); I’m trying to output the feature maps from a simple 1D convolutional model using a generator that collects signal data after the end of a 30-sec. epoch.

By performing a htop during the collection, I can clearly see memory usage is piling up and quickly causes OOM errors. This seems to happen at the forward phase of my network.

To make things understandable, here is a simplified version of my current code:

# encoding: utf-8
from torch.autograd import Variable
from functools import wraps
import progressbar
import numpy as np
import torch.nn as nn
import torch


class SimpleConvolution(nn.Module):
    def __init__(self, in_channels):
        super(SimpleConvolution, self).__init__()
        self.conv1 = nn.Conv1d(out_channels=8,
                               in_channels=in_channels,
                               kernel_size=4)

    def forward(self, x):
        return self.conv1(x)


def with_params_out(func):
    @wraps(func)
    def wrapper(*arg, **kwargs):
        generator = func(*arg, **kwargs)
        params = next(generator)
        return params, generator
    return wrapper


@with_params_out
def simple_generation(params):

    # Building a SimpleConvolution model.

    model = SimpleConvolution(params['number_of_channels'])
    if params.get('cuda', None):
        model = model.cuda()

    params_out = params.copy()
    epoch = int(params['fs'] * params['epoch'])

    # Initialize buffer & counters.

    signal_buffer = [0] * epoch
    signal_counter = 0

    cpt = 0
    signal = yield params_out

    while True:
        cpt += 1

        # Before reaching the end of an epoch,
        # collect information in the buffer.

        if signal_counter < epoch:
            signal_buffer[signal_counter] = signal
            signal_counter += 1

        # When reaching max limit, create a
        # Torch variable to contain it, and
        # forward data through the network.

        if signal_counter >= epoch:

            batch_eeg = Variable(
                torch.FloatTensor(np.array(signal_buffer)).view(1, -1, epoch),
                requires_grad=False,
                volatile=True)

        if cpt % epoch == 0:
            signal_counter = 0
            convolution_output = model.forward(batch_eeg)
            model.zero_grad()
            signal = yield convolution_output


if __name__ == "__main__":

    # Signal parameters

    params = {
        'fs': 250.,  # Sampling frequency
        'epoch': 30.,  # Epoch duration, in seconds
        'number_of_channels': 4,
        'cuda': False,
    }

    n_points = 10000
    inputs = {
        'data_max_size': n_points,
        'signal': np.random.uniform(-5e6, 5e6, size=(n_points, params['number_of_channels']))
    }

    # Generation procedure

    signal = []
    params, block = simple_generation(params)
    bar = progressbar.ProgressBar()

    for i in bar(range(inputs['data_max_size'])):

        result = block.send(tuple([inputs['signal'][i]]))
        signal.append(result[0].data.numpy())

    signal = np.array(signal)

I would appreciate if you could kindly suggest a workaround for this task to be performed with no memory leakage! Thank you very much in advance.

if you want to collect the output feature maps, collect their .data attributes. Remember that a Variable is not just it’s .data Tensor, but also all the history it’s holding onto (i.e. how it was created, what buffers it needs to compute gradients correctly, etc.).

I think every epoch when you yield and collect these params, it’s holding on to much more than just the output activations for this reason.