Pytorch: runtimeerror: mat1 dim 1 must match mat2 dim 0 using 2D CNN model with pretrained vgg_16

See more:

Python

artificial-intelligence

, +

I am working on cardiac CT data (axial, sagittal, coronal). I am using the pre-trained model vgg_16. But I got the following error. According to this error, my dimensions are not correct, but according to my code, I wrote things correctly. Can somebody guide me regarding this? I tried to correct my code already but got the same error. Below is the error and my code.

Traceback (most recent call last):
File "ct_pretrained.py", line 199, in <module>
loss, metric = train(model, train_loader, optimizer)
File "ct_pretrained.py", line 57, in train
output = model(axial, sagittal, coronal, emr)
File "/root/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "/data/heart_ct/torch/models/test.py", line 38, in forward
axial_feature = self.axial_model(axial)
File "/root/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "/root/miniconda/lib/python3.8/site-packages/torchvision/models/vgg.py", line 46, in forward
x = self.classifier(x)
File "/root/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "/root/miniconda/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 93, in forward
return F.linear(input, self.weight, self.bias)
File "/root/miniconda/lib/python3.8/site-packages/torch/nn/functional.py", line 1690, in linear
ret = torch.addmm(bias, input, weight.t())
RuntimeError: mat1 dim 1 must match mat2 dim 0

import torch
import torch.nn as nn
from torchvision import models

__all__ = ['VGG']

class VGG(nn.Module):

    def __init__(self, is_emr=False, mode='sum'):
        super().__init__()
        self.is_emr = is_emr
        self.mode = mode
        in_dim = 45

        self.axial_model = models.vgg16(pretrained=True)
        out_channels = self.axial_model.features[0].out_channels
        self.axial_model.features[0] = nn.Conv2d(1, out_channels, kernel_size=7, stride=1, padding=0, bias=False)
        self.axial_model.features[3] = nn.MaxPool2d(1)
        num_ftrs = self.axial_model.classifier.in_features #error in this line of code
        self.axial_model.classifier = nn.Linear(num_ftrs, 15)

        self.sa_co_model = models.vgg16(pretrained=True)
        self.sa_co_model.features[0] = nn.Conv2d(1, out_channels, kernel_size=7, stride=1, padding=(3,0), bias=False)
        self.sa_co_model.features[3] = nn.MaxPool2d(1)
        self.sa_co_model.classifier = nn.Linear(num_ftrs, 15)

        if self.is_emr:
            self.emr_model = EMRModel()
            if self.mode == 'concat': in_dim = 90

        self.classifier = Classifier(in_dim)

    def forward(self, axial, sagittal, coronal, emr):
        axial = axial[:,:,:-3,:-3]
        sagittal = sagittal[:,:,:,:-3]
        coronal = coronal[:,:,:,:-3]

        axial_feature = self.axial_model(axial)
        sagittal_feature = self.sa_co_model(sagittal)
        coronal_feature = self.sa_co_model(coronal)
        out = torch.cat([axial_feature, sagittal_feature, coronal_feature], dim=1)

        if self.is_emr:
            emr_feature = self.emr_model(emr)
            if self.mode == 'concat':
                out = torch.cat([out, emr_feature], dim=1)
            elif self.mode == 'sum':
                out += emr_feature

        out = self.classifier(out)

        return out

What I have tried:

I am using other models like ResNet/DenseNet it's working properly, but in vgg_16 it's getting errors in the code.