How to use TorchVisionFCModel to change model to accept 3D input?

[mattwarkentin]

Hi,

Based on the documentation it seems like monai.networks.nets.TorchVisionFCModel() allows one to use one of the existing torchvision models and adapt it so it accepts 3D data and can have any number of output classes. Maybe I am misunderstanding how this all works (it seems like it).

But I thought I could modify resnet50 to accept 3D data for a single class output like:

import monai

resnet50_3d = monai.networks.net.TorchVisionFCModel('resnet50', num_classes=1, dim=3)

The first layer of the network is now:

(pool): AvgPool3d(kernel_size=7, stride=1, padding=0)

But I can't seem to figure out what size the input tensor should be, and no matter what I've tried I get an error like the following:

resnet50_3d(torch.rand(1, 3, 224, 224, 224))

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/monai/networks/nets/netadapter.py", line 102, in forward
    x = self.features(x)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/torch/nn/modules/container.py", line 119, in forward
    input = module(input)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 399, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/Users/matt/.virtualenvs/torch/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 396, in _conv_forward
    self.padding, self.dilation, self.groups)
RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 7, 7], but got 5-dimensional input of size [1, 3, 224, 224, 224] instead

Am I doing this wrong?

[Nic-Ma]

Hi @mattwarkentin ,

Thanks for your interest here.
Sorry for the misleading document in TorchVisionFCModel, it can't adapt 3D data to 2D network.
The adapter means to support the 2D or 3D network with specified input classes.
I will try to enhance the document to make it much more clear.

Thanks.

[yiheng-wang-nv]

Hi @mattwarkentin , if you print the shape of the tensor before entering into the pooling layer in the forward function, you will see that it has the spatial shape (1, 1, 1):

import monai
import torchvision as tv
import torch

r3d = tv.models.video.r3d_18()
model = monai.networks.nets.NetAdapter(r3d, 1, 3, pool=("avg", {"kernel_size": 1, "stride": 1}))
output_before = model.features(torch.rand(1, 3, 112, 112, 112))
print(output_before.shape)

[yiheng-wang-nv]

Reference of the netadapter forward function: https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/netadapter.py#L94

[mattwarkentin]

Okay, my confusion was with how the layers are printed when you print the full NetAdapter instance. It shows the pool layer first, which made me think this was the first layer to receive the input tensor, even though the pool layer is called after model.features() in model.forward().

import monai
import torchvision as tv
import torch

r3d = tv.models.video.r3d_18()
model = monai.networks.nets.NetAdapter(r3d)

The order of calls is features -> pool [optional] -> flatten [optional] -> fc. But this is how it looks:

print(model)

NetAdapter(
  (pool): AvgPool2d(kernel_size=7, stride=1, padding=0)
  (features): Sequential(
    (0): BasicStem(
      (0): Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2), padding=(1, 3, 3), bias=False)
      (1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
    )
    (1): Sequential(
      (0): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
      )
      (1): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
      )
    )
    (2): Sequential(
      (0): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(64, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(2, 2, 2), bias=False)
          (1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
      )
    )
    (3): Sequential(
      (0): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(128, 256, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(2, 2, 2), bias=False)
          (1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
      )
    )
    (4): Sequential(
      (0): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(256, 512, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(512, 512, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv3d(256, 512, kernel_size=(1, 1, 1), stride=(2, 2, 2), bias=False)
          (1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): BasicBlock(
        (conv1): Sequential(
          (0): Conv3DSimple(512, 512, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (conv2): Sequential(
          (0): Conv3DSimple(512, 512, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)
          (1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (relu): ReLU(inplace=True)
      )
    )
    (5): AdaptiveAvgPool3d(output_size=(1, 1, 1))
  )
  (fc): Linear(in_features=512, out_features=1, bias=True)
)

I guess it would be nice if the __str__ was in the same order as forward.

[mattwarkentin]

I think flipping these lines around would fix the issue by adding features to the modules list before pool and fc:

MONAI/monai/networks/nets/netadapter.py

Lines 77 to 80 in 1c1c50f

	else:
	self.pool = get_pool_layer(name=pool, spatial_dims=dim)
	# remove the last 2 layers
	self.features = torch.nn.Sequential(*layers[:-2])

[yiheng-wang-nv]

Hi @mattwarkentin , thanks for your suggestion, and the implementation has been changed ( see #3228 )