How to load a Fastai model using Pytorch?

[ShubhankarRaj]

I have a fastai model was trained on RAVDESS's audio data to classify emotions:

from fastai.vision.augment import aug_transforms

tfms = aug_transforms(do_flip=False,flip_vert=False, max_lighting=0.1, max_zoom=1.05, max_warp=0.,max_rotate=5)

np.random.seed(42)

data = ImageDataLoaders.from_folder(processed_data_path, train=".", valid_pct=0.2, batch_tfms=[*tfms, Normalize.from_stats(*imagenet_stats)], size=144, num_workers=4)

learn = cnn_learner(data, models.resnet34, metrics=accuracy, lr=1e-4)

The remaining steps are the usual iterations of learn.fit_one_cycle where we are fine-tuning the model.
We created a fastai vision model using pretrained resnet34 model. The newly created model is loaded in Pytorch as below. NOTE that we are adding a Linear layer at the end of the fastai model:

import torch
import torch.nn as nn
from fastai.vision.models import resnet34
from fastai.vision.all import *
from fastai.vision.all import PILImage, tensor


class CustomModel(nn.Module):
    def __init__(self, pretrained_model_path):
        super(CustomModel, self).__init__()
        
        model_arch_path = f"{DATA_ROOT_PATH}/saved_models/emotion_detection_using_only_sound/fastai_model_arch.pth"
        base_model = torch.load(model_arch_path) #load model architecutre
        # base_model = resnet34()

        base_model.load_state_dict(torch.load(pretrained_model_path))

        # Freeze all layers
        for param in base_model.parameters():
            param.requires_grad = False

                
        # Get the last layer of the model
        last_layer = list(base_model.children())[-1]
        
        # Get the number of output features of the last layer
        if isinstance(last_layer, nn.Linear):
            num_features = last_layer.in_features
        else:
            num_features = last_layer[-1].in_features
        
        # Define a new linear layer with the same number of input features
        # and the desired number of output classes (e.g., 6 for 6 classes)
        self.new_layer = nn.Linear(num_features, 6)  # Assuming 6 classes for classification
        
        # Concatenate the base model with the new layer
        self.base_model = base_model
        # Replace the last layer with the new layer
        self.model = nn.Sequential(self.base_model, self.new_layer)


    def forward(self, x):
        # Convert JPEG file path to tensor
        image = PILImage.create(x)
        tensor_image = tensor(image)
        # Normalize image tensor
        normalized_image = tensor_image / 255.0
        # Add batch dimension
        normalized_image = normalized_image.unsqueeze(0)
        # Transpose to [batch_size, channels, height, width] as we were getting channel mismatch error
        normalized_image = normalized_image.permute(0, 3, 1, 2)

        # Move input tensor to GPU if available
        if torch.cuda.is_available():
            device = torch.device("cuda")
            normalized_image = normalized_image.to(device)
            self.to(device)
        
        print(normalized_image.shape)
        # Pass the normalized image through the base model
        features = self.model(normalized_image)

        print(features)
        # # Pass the features through the new linear layer
        logits = self.new_layer(features)
        return logits

WIth this code, the error that I am getting is : RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x6 and 512x6)

I understand that there is some issue in size compatibity of the model and input. But I am not able to figure out how to resolve it.