预训练模型加载&更改网络层&batch推断

最新推荐文章于 2024-07-17 00:32:55 发布

~~小智~~

最新推荐文章于 2024-07-17 00:32:55 发布

阅读量433

点赞数

文章标签： python 深度学习

本文链接：https://blog.csdn.net/daydayup_study/article/details/116528754

版权

该博客介绍了如何使用预训练的MobileNetV3模型进行迁移学习，步骤包括加载预训练模型，修改网络的最后一层以适应不同类别数，固定部分层以进行微调，以及如何进行推断时的图像处理。此外，还提供了测试代码和如何在项目中固定和重组模型层的细节。

摘要由CSDN通过智能技术生成

以MobileNetV3为例，预训练模型classes=1000，项目中classes=2, 利用预训练模型训练自己的任务。

1. 预训练模型加载

net = MobileNetV3_Small(num_classes=1000)
##加载预训练模型
model = torch.load('./mbv3_small.pth.tar', map_location = 'cpu')
print('Loading base network...')
weight = model["state_dict"]
base_net = torch.nn.DataParallel(net)
base_net.load_state_dict(weight)

注意 base_net = torch.nn.DataParallel(net)，那么推断的时候也须加上：

model = torch.nn.DataParallel(model)

2. 更改网络层

将最后一层全连接层替换：

base_net.module.linear4 = nn.Linear(1280, 2).cuda()

3. 测试：

x = torch.randn(2,3,224,224).cuda()
y = base_net(x)
print(y.size())

4. 有些时候要固定某些层，方法如下：

#固定除了linear4层外的其他层
for k,v in base_net.named_parameters():
    if k!='linear4.weight' and k!='linear4.bias' :
        v.requires_grad=False#固定参数

参考：

https://www.jianshu.com/p/d67d62982a24

https://blog.csdn.net/qq_28368377/article/details/108747997

https://blog.csdn.net/qq_43360777/article/details/106305469?utm_medium=distribute.pc_relevant.none-task-blog-baidujs_title-0&spm=1001.2101.3001.4242

5. 有时候推断的时候需要将图片重组为batch，方法如下：

def inference_detectors(model,frames):
    batch_frame = []
    for frame in frames:
        frame = np.asarray(frame)
        ##完成按长宽比对其进行缩放
        border_v = 0
        border_h = 0
        if (frame.shape[0] / frame.shape[1]) < 1:
            border_v = int(((frame.shape[1]) - frame.shape[0]) / 2)
        else:
            border_h = int(((frame.shape[0]) - frame.shape[1]) / 2)
        img = cv2.copyMakeBorder(frame, border_v, border_v, border_h, border_h, cv2.BORDER_CONSTANT, 0)
        img = cv2.resize(img, (224, 224))
        #img = cv2.resize(np.asarray(frame),(224,224))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # TORGB
        img = transforms.ToTensor()(img)
        batch_frame.append(img)

    if len(batch_frame) == 0:
        return []

    ##转换为batch
    batch_frame = torch.stack(batch_frame,dim=0)
    img = Variable(torch.tensor(batch_frame).float(), requires_grad=False).cuda()
    output = model(img)
   

    pred = output.argmax(dim=1)
    pred_cpu = pred.data.cpu()
    index = np.argwhere(pred_cpu == 1)
    
    result = []
    ##返回index是一个list, 取值方法如下：
    if len(indexs[0]) > 0:
       for id in indexs[0]:
        bbox = bboxes[id]
        result.append([bbox[0], bbox[1], bbox[2], bbox[3]])    
    return result