Implement a 3D custom model

Hi everyone, I would like to implement a 3D, single channel, PyTorch model having two 3D convolutions, followed by two linear layers, independently from the batch size and shape of the input volume.

Any suggestions?


Something like this is similar to what you’re asking:

Except that has more than 2 channels and some added pooling/dropout/batchnorm.

Should help you on the right track.

Thank you. Only one question what are sample_duration and sample_size ?

Not my code, but that is used for calculating the transition between the final conv layer and the first Linear layer. However, you can just use nn.LazyLinear:

Just specify the out_features you want and it will figure out the input size during the first run.

Thanks a lot @J_Johnson !
The LazyLinear was really useful.

I wrote this snippet of code (batch = 2, 1 channel, 130x130x130 image) and I expected to get a tensor having shape (2,3), but I got a (2,64,64,64,3).

What am I missing?
Thanks a lot.

import torch
import torch.nn as nn
import torch.nn.functional as F

class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()
        self.seq1 = nn.Sequential(
            nn.Conv3d(in_channels=1, out_channels=64, kernel_size=3, stride=1),
        self.l1 = nn.LazyLinear(out_features=64, bias=False)
        self.l2 = nn.Linear(64, 3, bias=False)

    def forward(self, x):
        x1 = self.seq1(x)
        x2 = self.l1(x1)
        x3 = self.l2(x2)

        return x3

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
fake_input = torch.randn((2,1,130,130,130)).to(device) # Batch = 2, 1 channel, 130x130x130 image

model = CustomNet().to(device)

outputs = model(fake_input)


This post has the output size calculation from the docs written as a Python function. You can enter in the arguments on one dimension for your Conv3d and MaxPool3d as well as input size on one dimension to determine what the output size is you should expect from that layer, given a certain input size.

Thanks @J_Johnson, I really appreciate your reply. I solved by adding a Flatten before the LazyLinear layer to my custom model.

Have a nice day.