how to use my custom trained model

google-research / big_transfer

Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.

https://arxiv.org/abs/1912.11370

Apache License 2.0

1.5k stars 175 forks source link

how to use my custom trained model #51

Open kyuhong0206 opened 3 years ago

kyuhong0206 commented 3 years ago

i finetuned BiT-R50 and save like bit.pth.tar torch.save({ "step": step, "model": model.state_dict(), "optim" : optim.state_dict(), }, savename)

but i can't use it for test like checkpoint = torch.load(f'/srv/data/datasets/kyuhong/big_transfer/bit_logs/products/bit.pth.tar') self.model.load_state_dict(checkpoint['model'])

how can i use my custom trained model