DeepSets with pre-trained image model


I am fairly new with fastai. I am trying to implement the DeepSets paper which basically takes multiples images as input where the order of the images don’t matter. For example takes multiples MNIST images and return the sum of the numbers in the images.

Basically you use a reset for each image and then use summation on the output vectors of the resnet for the images and then a fully connected layer for the regression task of calculating the summation.

I wanted to use pre-trained reset34 for that to leverage transfer learning.

I guess I will need a new ItemBase for an N-images input (I will adapt the ImageTuple).

I found an implementation in Python of DeepSets:

But my main question is around using the same pre-trained model (without a custom head) on multiple images and then summing their ouput vector. Basically I would like to replace SmallMNISTCNNPhi from the code example above by a pre-trained model.

Also is there any way I could still use freeze/unfreeze for the pre-trained model in this situation.