import torch import torchvision from torch import nn # ==================== 1. 加载预训练模型 ==================== print("=" * 50) print("1. 加载预训练 VGG16 模型") vgg16 = torchvision.models.vgg16(pretrained=True) 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(pretrained=False) 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(pretrained=True)
pretrained=True 模型vgg16已经拿超级多的图片训练过了=一个见多识广的老人
pretrained=Flase 模型vgg16没训练过=一个小孩子
2.原来是Linear(in_features=4096, out_features=1000, bias=True)
修改有两种方式，方式一直接修改：vgg16.classifier[6] = nn.Linear(4096, 10)
Linear(in_features=4096, out_features=10, bias=True)
方式二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(pretrained=False) vgg16_reload.classifier[6] = nn.Linear(4096, 10) # 先做同样的修改 vgg16_reload.load_state_dict(torch.load("vgg16_method2.pth"))
两种保存实质保存的是一个内容，所以卷积核一样
5.为什么要保存pth文件，直接用python文件不好吗？
python文件保存的是目录，不是内容，不是具体的卷积核参数，你保存python文件下一次还得重新加载。保存pth模型文件，记录了上次训练的结果，下次拿来可以直接用
