When I try to train my model which contains MaxPool3d , It always end up with ‘out of memory’ error.
my environment info is here:
Ubuntu 16.04
PyTorch version: 0.4.0a0+da6c3c9
installed PyTorch from source
python version is 3.6.2
CUDA/cuDNN version: 9.1/7.1.2
GCC version (if compiling from source):GCC 5.4
Build command you used (if compiling from source): as default
I can reproduce this bug by the following script:
import torch
import torch.nn as nn
from torch import optim
from torch.nn.functional import smooth_l1_loss
model = nn.Sequential(
nn.Conv3d(1,1,kernel_size=3,padding=1),
nn.ReLU(),
nn.MaxPool3d(kernel_size=3, stride=2, padding=1),
)
#optimizer = optim.SGD((x for x in model.parameters() if x.requires_grad is True), lr=1e-3, momentum=0.9,nesterov=True)
crit = smooth_l1_loss
model.cuda()
count = 0
while True:
# optimizer.zero_grad()
input = torch.rand(30,1,200,200,200).cuda()
loc_output= model(input)
loc_outpus = loc_output
if type(loc_output) == tuple:
loc_output = loc_output[0]
targets = torch.rand(loc_output.size()).cuda()
loss = crit(loc_output,targets)
loss.backward()
# optimizer.step()
del loss,loc_output,targets,input
torch.cuda.empty_cache()
count += 1
print(count)