So @wyli, This patch works on models like Inception_V3 and produces the desired output. However, it fails for models like VisionTransformer torchnvision.models.vit_b_16 with the following error

  File "/raid/Vikash/Tools/review-MONAI/MONAI/monai/networks/nets/torchvision_fc.py", line 91, in __init__
    raise ValueError(f"Model ['{model_name}'] does not have a Linear layer at the end.")
ValueError: Model ['vit_b_16'] does not have a Linear layer at the end.

The error is obvious as there is the startswith Linear condition check. It might work for Inception and Resnet as the tail of the model looks like this

  )
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (dropout): Dropout(p=0.5, inplace=False)
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)

But does not work if the model terminates in a Linear layer wrapped in a Sequential module as follows

          (4): Dropout(p=0.0, inplace=False)
        )
      )
    )
    (ln): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
  )
  (heads): Sequential(
    (head): Linear(in_features=768, out_features=1000, bias=True)
  )
)

So, if we can expand the if statement it will cover majority of the models in torchvision.models

thanks @vikashg the comments are helpful, I spend some time on this, it's not easy to be that flexible at the same time backwards compatible... but for the basic use cases it's ok now:

import torch
from torchvision.models import Swin_T_Weights

from monai.networks.nets import TorchVisionFCModel

model = TorchVisionFCModel(
    "vit_b_16", num_classes=4, use_conv=False, pool=None, in_channels=768, fc_name="heads.head")
output = model.forward(torch.randn(2, 3, 224, 224))
print(output.shape)  # torch.Size([2, 4])

model = TorchVisionFCModel(
    "swin_t",
    num_classes=4,
    pool="",
    use_conv=True,
    in_channels=768,
    node_name="permute",
    weights=Swin_T_Weights.IMAGENET1K_V1,
)
output = model.forward(torch.randn(2, 3, 500, 500))
print(output.shape)  # torch.Size([2, 4, 16, 16])

/black

Thanks @wyli for supplying an example code snippet. It all works well now. IT should be compatible with 0.9 right ?

yes it's non-breaking, the previous test cases all work fine.

/build

/build

/build

/build

/build

/build

Will these changes be available if we install monai via pip 0.8.1? Because we are using NVFlare for federated learning following this tutorial which uses pip install monai=0.8.1 to build the docker packages. I wonder if these changes will be pushed to the pip package manager.

Hi @ctestagrose, there's currently no plan to backport this to v0.8, it will be available in the weekly dev release https://pypi.org/project/monai-weekly/#history (next one on 14th/Aug). So, early next week, you can uninstall monai with pip uninstall -y monai and run pip install monai-weekly to get this feature. This version will be very different from v0.8.1 0.8.1...dev including some breaking changes so the existing 0.8.1 pipelines may or may not work.