从vgg16-397923af.pth里读取的数值应该和加载预训练模型后model.load_state_dict参数一致

深浅卡布星

于 2023-07-17 19:31:37 发布

阅读量487

点赞数

分类专栏： Experiment 文章标签： pytorch

本文链接：https://blog.csdn.net/hahhahahhaja/article/details/131773400

版权

Experiment 专栏收录该内容

11 篇文章 0 订阅

订阅专栏

import numpy as np
import torch
import torchvision
from torchvision import models
from PIL import Image
from matplotlib import pyplot as plt
import torchvision.transforms as transforms
from src import fcn_resnet50, resnet50
from model_fcn8s import VGG, fcn_vgg16

pretrain_backbone = True
struct = [(2, 64), (2, 128), (3, 256), (3, 512), (3, 512)]
backbone = VGG(num_classes=21, struct=struct)
# vgg16 = models.vgg16(pretrained=False)
if pretrain_backbone is True:
    weights_dict = torch.load("vgg16-397923af.pth", map_location='cpu')
    for name, param in weights_dict.items():
        print(f"{name}: {param}")

    backbone.load_state_dict(weights_dict, strict=False)
    print("Loaded pretrained model parameters:")
    for name, param in backbone.state_dict().items():
        print(f"{name}: {param}")