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.