pytorch15-＞vgg16的修改，保存，加载
import torch import torchvision from torch import nn # 1. 加载预训练模型 print( * 50) print(1. 加载预训练 VGG16 模型) vgg16 torchvision.models.vgg16(pretrainedTrue) print(vgg16) # 会打印出完整的网络结构 # 2. 修改模型 (适配 CIFAR-10) print(\n * 50) print(2. 修改最后一层为 10 分类 (适配 CIFAR-10)) # 方法1直接替换推荐 vgg16.classifier[6] nn.Linear(4096, 10) # 方法2追加一层 # vgg16.classifier.add_module(my_linear, nn.Linear(1000, 10)) print(vgg16) # 打印修改后的结构 # 3. 模型保存 (P26) print(\n * 50) print(3. 保存模型) # 方式1保存完整模型结构 参数 torch.save(vgg16, vgg16_method1.pth) # 方式2只保存参数推荐 torch.save(vgg16.state_dict(), vgg16_method2.pth) print(模型已保存) # 4. 模型加载 (P26) print(\n * 50) print(4. 加载模型) # 加载方式1加载完整模型 # 注意这样加载需要模型类的定义在当前文件中如果是自定义模型需要先定义类 model1 torch.load(vgg16_method1.pth) print(方式1加载成功) # 加载方式2先重建模型结构再加载参数推荐 # 注意重建时必须与保存时的结构完全一致 vgg16_reload torchvision.models.vgg16(pretrainedFalse) vgg16_reload.classifier[6] nn.Linear(4096, 10) # 先做同样的修改 vgg16_reload.load_state_dict(torch.load(vgg16_method2.pth)) print(方式2加载成功) # 5. 验证参数一致 print(\n * 50) print(5. 验证参数是否一致) # 比较两种方式加载的模型的第一层卷积核是否相同 same torch.equal(model1.features[0].weight, vgg16_reload.features[0].weight) print(f两种方式加载的模型参数是否一致 {same}) print(\n✅ 所有操作完成)1.vgg16 torchvision.models.vgg16(pretrainedTrue)pretrainedTrue 模型vgg16已经拿超级多的图片训练过了一个见多识广的老人pretrainedFlase 模型vgg16没训练过一个小孩子2.原来是Linear(in_features4096, out_features1000, biasTrue)修改有两种方式方式一直接修改vgg16.classifier[6] nn.Linear(4096, 10)Linear(in_features4096, out_features10, biasTrue)方式二vgg16.classifier.add_module(my_linear, nn.Linear(1000, 10))这是打补丁推荐使用方式一。3.保存有两种方式方式一保存结构和参数torch.save(vgg16, vgg16_method1.pth)方式二用字典只保存参数torch.save(vgg16.state_dict(), vgg16_method2.pth)相当于你要去图书馆看书方式一是你把连书带书架全都借走了方式二是只借书推荐方式二4.但是使用方式二你再次加载时需要重新加载骨架因为你保存时就没保存骨架像这样vgg16_reload torchvision.models.vgg16(pretrainedFalse) vgg16_reload.classifier[6] nn.Linear(4096, 10) # 先做同样的修改 vgg16_reload.load_state_dict(torch.load(vgg16_method2.pth))两种保存实质保存的是一个内容所以卷积核一样5.为什么要保存pth文件直接用python文件不好吗python文件保存的是目录不是内容不是具体的卷积核参数你保存python文件下一次还得重新加载。保存pth模型文件记录了上次训练的结果下次拿来可以直接用
