How to transfer the trained ckpt to pytorch.pth model?

[knaffe]

I use the example script to train a model, I got a ckpt file. but how could I extra the trained resnet50.pth instead of the whole SelfSupervisedLearner? Sorry I am new for pytorch lightning lib.
What I want is the SelfSupervised resnet50.pth, because I want this to replace the original ImageNet-pretrained one.
Thank you a lot.

[lucidrains]

@knaffe Oops, yea, I realized there isn't an easy way to access the original net

Try updating to 0.5.2, and then, from your SelfSupervisedLearner instance, do resnet = model.learner.net

[knaffe]

ok. I think your implt would work.
But if I do this below:

resnet = models.resnet50(pretrained=True)   
model = SelfSupervisedLearner(resnet, xxx)  
trainer = pl.Trainer(xxx)  
trainer.fit(model, train_loader)  
torch.save(resnet.state_dict(), './improved-resnet50.pt')  
 

When finishing training, could I save the correct trained resnet-50.pth？

[lucidrains]

@knaffe yup, that works too!

[yangyangkiki]

Hi all,
Thank you for sharing your codes!
I was wondering where I can download the imagenet pretrained resnet50 model by using BYOL?
Thank you so much!

[utkutpcgl]

Hi, thanks for the great repo.
Pretrained models would be great as @yangyangkiki said.