Pytorch中只导入部分层权重的方法

我们通常会用到迁移学习,即在一个比较通用的pretext-task上做预训练,随后针对不同的downstream task进行微调。而在微调的时候,网络结构的最后几层通常是要做出改变的。举个例子,假设pretext-task是在imagenet上面做图像分类,而下游任务是做语义分割,那么在微调的时候需要将分类网络的最后几层全连接层去掉,改造成FCN的网络结构。此时就需要我们把前面层的权重加载进去。

如果改了模型结构以后,再简单粗暴的使用torch.load_state_dict(torch.load(‘xxx.pth’))那么肯定就会报错。所以具体怎么办呢,且耐心往下看。

首先我们定义一个简单的图像分类模型:

class model1(nn.Module):
    def __init__(self, img_size):
        super(model, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, 1, 1)
        self.conv2 = nn.Conv2d(16, 64, 3, 1, 1)
        self.fc1 = nn.Linear(self.num_feature_pixel(img_size), 1024)
        self.fc2 = nn.Linear(1024, 10)

    def
评论 5
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值