How to save part of model parameters?


(손주형) #1

I’m using vgg16 backbone on my mode and I want to save just this fine-tuned vgg16 parameters.
how can i do that?


(Anton) #2

Hi, you can do smth like

torch.save({
        "epoch": epoch + 1,
        "model_state_dict": model.state_dict(),
        ...
        "optimizer": optimizer.state_dict(),
        "loss": loss
    }, model_path)

but specify only parameters you want to save


(손주형) #3

how can i specify parameters that i want?


(Anton) #4

Just pass to save the parameters you want. E.g.

...
"my_param": parameters,
...