Я пытаюсь обучить нейронную сеть для очень большого ввода (5 * 100 000 000), и это требует гораздо больше памяти, чем ожидалось.Вот некоторый минимальный пример:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import time
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv1d(in_channels=5, out_channels=1, kernel_size=100000000, stride=10)
def forward(self, x):
x = self.conv1(x)
x = torch.sigmoid(x)
return x
model = Net().cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.BCELoss()
data = torch.normal(torch.zeros(1,5,100000000),torch.ones(1,5,100000000))
data = data.cuda()
label = torch.ones(1,1,1)
label = label.cuda()
for epoch in range(10):
output = model(data)
loss = criterion(output, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Epoch :", epoch)
На входе находятся некоторые случайные данные, они используют приблизительно 2 ГБ, как и ожидалось (32 бита * 5 * 100 000 000 = 1,86 ГБ). Эта переменная не имеет градиента.Сеть состоит из одного сверточного слоя с одним фильтром того же размера, что и вход, поэтому она имеет вес 500 М, то есть еще 2 ГБ.После форварда пас еще 2Гб привыкаешь.После использования loss.backprop() 8 ГБ, после использования optimizer.step() 12 ГБ, то есть всей доступной памяти.
Во время второго прохода в эпоху идет нормально, но во время обратного распространения я получаю RuntimeError: ошибка CUDA: outпамяти.
Что именно сохранялось в памяти GPU в эпоху?Почему память не освобождается после завершения шага оптимизации?Как уменьшить использование памяти в этом случае?
UPD: похоже, мои проблемы похожи на эту проблему https://discuss.pytorch.org/t/how-to-free-gpu-memory-and-delete-memory-allocated-variables/20856