I have a Pytorch Lightning application that I use to train a CNN for image classification from scratch using NNI. I later use the trained model in production to do predictions.
As transfer learning seems to be much simpler and faster with fastai, I would like to migrate the training component to fastai while keeping the old prediction code as it is for now (Pytorch Lightning). I already used a vision_learner to do that and exported the resulting model as .pkl (.pth) file.
In theory, it should be straight-forward to load the model in Pytorch:
Your LightningModule is completely empty. There is nothing you can load the state_dict into and it has no forward method. There is no way to infer the model architecture from only the state dict. The state dict only includes the layers weights.
What network architecture did you use in your fastai model? I would create a LightningModule with the architecture in it and define a forward method, and then the second approach with the .pth file should work.
Ok, I had to call .float() on my tensor to fix the above mentioned error.
representations = self.feature_extractor(x.float()).flatten(1)
Now that I got this working, I discovered that ResNet-50 is trained with images of 224 pixels x 224 pixels. While I can fine_tune the model with images with other dimensions in FastAI, it seems not to work in the prediction stage if I understand the error message correctly.
RuntimeError: mat1 and mat2 shapes cannot be multiplied (262144x1 and 2048x2)