Hi,
I after struggeling for a day im stuck, any ideas would be greatly appriciated!
I’m trying to do a 3D convolution on a synthetic MNIST dataset where the images are stretched in a third dimension with torch.einsum('i,jkl->jikl', torch.ones(28), img) s.t. the dimension become (1,1,28,28,28).
The following code in VSCode gives me a roughly 20% chance of nan value in the lss.backward() calculation over a single image.
import numpy as np  
import torch
from torch import nn, optim
import torchvision
import random

torch.cuda.manual_seed(5)
torch.manual_seed(5)
np.random.seed(42)
random.seed(42)


class AddDimension(torch.nn.Module):
    
    def __init__(self):
        super(AddDimension, self).__init__()
        
    def forward(self, img):
        return torch.einsum('i,jkl->jikl', torch.ones(28), img)
        


transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),
    AddDimension()
])

train_dataset = torchvision.datasets.MNIST(
    root="~/torch_datasets", train=True, transform=transform, download=True
)

test_dataset = torchvision.datasets.MNIST(
    root="~/torch_datasets", train=False, transform=transform, download=True
)

train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True
)

test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=1, shuffle=False, num_workers=4
)



class ConvAutoEncoder4(torch.nn.Module):
    def __init__(self):
        super(ConvAutoEncoder4, self).__init__()
        channels = 4
        self.fcn3 = nn.Conv3d(in_channels=1, out_channels=channels, kernel_size=4, stride=4)
        self.fcn4 = nn.Conv3d(in_channels=channels, out_channels=channels, kernel_size=1)
        self.fcn5 = nn.Conv3d(in_channels=channels, out_channels=channels, kernel_size=1)
        self.fcn6 = nn.ConvTranspose3d(in_channels=channels, out_channels=1, kernel_size=4, stride=4)
        self.print = False
        
    def set_print(self, pr):
        self.print = pr
        
    def forward(self, tensor):
        if not self.print:
            tensor = self.fcn3(tensor)
            tensor = torch.relu(tensor)
            tensor = self.fcn4(tensor)
            tensor = torch.relu(tensor)
            tensor = self.fcn5(tensor)
            tensor = torch.relu(tensor)
            tensor = self.fcn6(tensor)
            tensor = torch.relu(tensor)
            return tensor
        
        tensor = self.fcn3(tensor)
        tensor = torch.relu(tensor)
        print(tensor.shape)
#         for slc in tensor[0][0]:
#             print(slc)
#         for name, param in self.named_parameters():
#             print(name, torch.isfinite(param.grad).all())
        tensor = self.fcn4(tensor) 
        tensor = torch.relu(tensor)
        tensor = self.fcn5(tensor)
        tensor = torch.relu(tensor)
        tensor = self.fcn6(tensor)
        tensor = torch.relu(tensor)
        for name, param in self.named_parameters():
            print(name, param.grad)
        return tensor


it = iter(train_loader)
imgs = [next(it)[0] for i in range(100)]

print(len(imgs))

print(imgs[0].shape)

mod = ConvAutoEncoder4()
crtrn = nn.MSELoss()
ptmzr = optim.Adam(mod.parameters(), lr=1e-3)

torch.autograd.set_detect_anomaly(True)

mod.set_print(True)

ptmzr.zero_grad()
res = mod(imgs[0])
lss = crtrn(res, imgs[0])
lss.backward()
ptmzr.step()
print(lss.item())


With the following stacktrace:
  File "/Users/.../python/MNIST3D.py", line 118, in <module>
    res = mod(imgs[0])
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/.../python/MNIST3D.py", line 93, in forward
    tensor = self.fcn5(tensor)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 572, in forward
    return F.conv3d(input, self.weight, self.bias, self.stride,
 (Triggered internally at  ../torch/csrc/autograd/python_anomaly_mode.cpp:104.)
  Variable._execution_engine.run_backward(
Traceback (most recent call last):
  File "/Users/.../python/MNIST3D.py", line 120, in <module>
    lss.backward()
  File "/usr/local/lib/python3.8/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/usr/local/lib/python3.8/site-packages/torch/autograd/__init__.py", line 130, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Function 'SlowConv3DBackward' returned nan values in its 1th output.

Any clues?