I am implementing an encoder-decoder network in which the encoder takes a 5d input and compresses it to a 4d output, while the decoder takes a 4d input and up-samples to a 4d output. There are shortcut connections which pass 4d slices of encoder feature maps to the decoder. There are also 4d slices of MaxPool indices passed from encoder to decoder. I am getting “illegal memory access” errors. Here is a minimal example:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torch.utils.data
#---------------------------------------------------------------------------
# define encoder-decoder network and optimizer
class encoder_decoder(nn.Module):
def __init__(self):
super().__init__()
self.enc_conv1 = nn.Conv3d(1,8, kernel_size=(1,3,3), stride=1, padding=(0,1,1))
self.pool = nn.MaxPool3d(kernel_size=(1,2,2), stride=(1,2,2), return_indices=True)
self.enc_conv2 = nn.Conv3d(8,8, kernel_size=(3,1,1), stride=1, padding=0)
self.dec_conv1 = nn.Conv2d(8,8, kernel_size=3, stride=1, padding=1)
self.unpool = nn.MaxUnpool2d(kernel_size=2, stride=2)
self.dec_conv2 = nn.Conv2d(16,1, kernel_size=3, stride=1, padding=1)
def forward(self, net_input):
out = net_input
out = self.enc_conv1(out) # 1x3x64x64 -> 8x3x64x64
shortcut = out[:,:,1] # reference to center slice of out
out, indices = self.pool(out) # 8x3x64x64 -> 8x3x32x32
indices = indices[:,:,1] # reference to center slice of indices
out = self.enc_conv2(out) # 8x3x32x32 -> 8x1x32x32
out = out.squeeze(2) # 8x1x32x32 -> 8x32x32
out = self.dec_conv1(out) # 8x32x32 -> 8x32x32
out = self.unpool(out, indices) # 8x32x32 -> 8x64x64
out = torch.cat((shortcut,out),1) # 8x64x64 -> 16x64x64
out = self.dec_conv2(out) # 16x64x64 -> 1x64x64
return out
net = encoder_decoder()
net.cuda()
criterion = nn.MSELoss()
criterion.cuda()
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.5)
#----------------------------------------------------------------------------
# define dataset and dataloader
class create_dataset(torch.utils.data.Dataset):
def __init__(self):
self.data = [ (torch.rand(1,3,64,64),
torch.rand(1,64,64)) for i in range(100) ]
def __getitem__(self, idx):
return self.data[idx]
def __len__(self):
return len(self.data)
dataset = create_dataset()
print('Loaded ' + str(len(dataset)) + ' training examples')
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=True, num_workers=8)
#----------------------------------------------------------------------------
# training loop
print('Start training loop')
for epoch in range(4):
print('Epoch: ' + str(epoch))
net.train()
for training_idx, (input_batch,target_batch) in enumerate(dataloader):
print('Training batch: ' + str(training_idx))
input_batch = Variable(input_batch.cuda())
target_batch = Variable(target_batch.cuda())
optimizer.zero_grad()
output_batch = net(input_batch)
err = criterion(output_batch, target_batch)
err.backward()
optimizer.step()
Here is the output:
$ CUDA_LAUNCH_BLOCKING=1 python error.py
Loaded 100 training examples
Start training loop
Epoch: 0
Training batch: 0
Training batch: 1
THCudaCheck FAIL file=/py/conda-bld/pytorch_1490895093647/work/torch/lib/THCUNN/generic/SpatialMaxUnpooling.cu line=43 error=77 : an illegal memory access was encountered
Traceback (most recent call last):
File "error.py", line 78, in <module>
output_batch = net(input_batch)
File "/home/ubuntu/anaconda3/lib/python3.5/site-packages/torch/nn/modules/module.py", line 206, in __call__
result = self.forward(*input, **kwargs)
File "error.py", line 31, in forward
out = self.unpool(out, indices) # 8x32x32 -> 8x64x64
File "/home/ubuntu/anaconda3/lib/python3.5/site-packages/torch/nn/modules/module.py", line 206, in __call__
result = self.forward(*input, **kwargs)
File "/home/ubuntu/anaconda3/lib/python3.5/site-packages/torch/nn/modules/pooling.py", line 304, in forward
self.padding, output_size)
File "/home/ubuntu/anaconda3/lib/python3.5/site-packages/torch/nn/functional.py", line 277, in max_unpool2d
return f(input, indices)
File "/home/ubuntu/anaconda3/lib/python3.5/site-packages/torch/nn/_functions/thnn/pooling.py", line 177, in forward
self.output_size[1], self.output_size[0])
RuntimeError: cuda runtime error (77) : an illegal memory access was encountered at /py/conda-bld/pytorch_1490895093647/work/torch/lib/THCUNN/generic/SpatialMaxUnpooling.cu:43