🔥码云GVP开源项目 12k star Uniapp+ElementUI 功能强大 支持多语言、二开方便! 广告
```py # 1. Directly Load a Pre-trained Model # https://github.com/pytorch/vision/tree/master/torchvision/models import torchvision.models as models resnet50 = models.resnet50(pretrained=True) # or model = models.resnet50(pretrained=False) # Maybe you want to modify the last fc layer? resnet.fc = nn.Linear(2048, 2) # 2. Load part of parameters of a pretrained model as init for self-defined similar-architecture model. # resnet50 is a pretrain model # self_defined indicates model you just define. resnet50 = models.resnet50(pretrained=True) self_defined = ... pretrained_dict = resnet50.state_dict() model_dict = self_defined.state_dict() pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} # update & load model_dict.update(pretrained_dict) model.load_state_dict(model_dict) # 3. Save & Load routines. # routine # (recommended) saves and loads only the model parameters torch.save(model.state_dict(), PATH) model.load_state_dict(torch.load(PATH)) # routine 2 # saves and loads the entire model torch.save(model, PATH) model = torch.load(PATH) ```