Classification backbone with Vit results in argument 'input' (position1) must be Tensor, not tuple

[kavmar]

Hi,

I am trying to use ViT as follows:

net = monai.networks.nets.ViT(spatial_dims=2, in_channels=1, img_size=(400, 400), proj_type='conv', patch_size=(64, 64),
num_classes=6, classification=True, post_activation='0').to(device)

but I am running into the same issue as reported here: #464

return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
TypeError: cross_entropy_loss(): argument 'input' (position 1) must be Tensor, not tuple

It has been concluded that the API will be enhanced by hidden_states_out, but I do not see it implemented - apparently due to design.

MONAI version: 1.3.0
Pytorch version: 2.1.1+cu121

Thanks for advice

[KumoLiu]

Hi @kavmar, I think you can take outputs[0] for the loss instead of just outputs.
#464 (comment)

Hope it helps, thanks!

[kavmar]

Hi @KumoLiu , and thanks for suggestion. I saw this workaround in the associated discussions. Sorry I forgot to mention that I am using monai engine SupervisedEvaluator/SupervisedTrainer, where I do not have outptuts from loss in immediate reach. Any suggestion how to tackle this in this situation?

Thanks a lot

[ericspod]

The issue is that the network puts out a tuple of values and you only want the first one, there's a few ways to create a ad-hoc solution. You can inherit from ViT with only a __call__ method which returns the first value from the parent call:

class MyViT(ViT):
    def __call__(self,*args,**kwargs): return super().__call__(*args,**kwargs)[0]

You could instead define an Inferer subclass to do the equivalent thing:

class MySimpleInferer(SimpleInferer):
    def __call__(self,*args,**kwargs): return super().__call__(*args,**kwargs)[0]

You would use this in place of the default SimpleInferer class that the trainer/evaluator classes use. Both these throw out the hidden states so if you want these in the state object this won't do.

You could instead use a lambda to create a new loss function which passes along only the first value of the target argument:

def my_loss(input,target):
    return original_loss(input, target[0])

This will keep the hidden states in the state object and only use the first tensor in the tuple as you want. This would be the loss function to pass to the engines in place of original_loss. I guess you can experiment and see which of these solutions work.

[kavmar]

Wonderful, @ericspod. Many thanks for the copy and paste solution. I decided to go MyViT, because it has the least effect on the rest of the code.