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

Saving model after finetuning #33

Open sahilwalke17 opened 4 years ago

sahilwalke17 commented 4 years ago

How do I save the model after fine-tuning using command line on dataset(CIFAR100) ?

sickerin commented 4 years ago

you gotta edit a line in train.py. In the PyTorch scripts, apparently the block that calls torch.save only gets executed when you enable the eval_every flag.

you can add this line torch.save at the end of the training so that it saves one final time before it finishes training

    # Final eval at end of training. and save model?
    run_eval(model, valid_loader, device, chrono, logger, step='end')
    torch.save({
                "step": step,
                "model": model.state_dict(),
                "optim" : optim.state_dict(),
            }, savename)
charliewang789 commented 3 years ago

thx. this saved as a tar file, which contains a set of files. should it be a simple pt file?

tianrengao commented 3 years ago

I did as skcerin said. But failed to load to train, it says "expected cpu but got cuda:0"

kyuhong0206 commented 3 years ago

you gotta edit a line in train.py. In the PyTorch scripts, apparently the block that calls torch.save only gets executed when you enable the eval_every flag.

you can add this line torch.save at the end of the training so that it saves one final time before it finishes training

    # Final eval at end of training. and save model?
    run_eval(model, valid_loader, device, chrono, logger, step='end')
    torch.save({
                "step": step,
                "model": model.state_dict(),
                "optim" : optim.state_dict(),
            }, savename)

how can use it? for test