智能数字图像处理：MobileNetV2代码（pytorch）之train.py解读

最新推荐文章于 2024-02-04 16:51:45 发布

原创

最新推荐文章于 2024-02-04 16:51:45 发布 · 1.7k 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #cuda #processing

本文介绍了使用PyTorch实现MobileNetV2模型进行图像分类的过程，包括模型实例化、加载预训练权重、筛选权重字典、加载权重、冻结特征提取层的权重等关键步骤。

1.net = MobileNetV2(num_classes=5)-》实例化模型

2.torch.load()函数载入预训练模型。

3.pre_dict = {k: v for k, v in pre_weights.items() if "classifier" not in k}-》遍历权重字典，看是否有classifier整个参数，如果不在层名称当中则进行一个保存。

4.missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict=False)-》载入权重字典。

5.for param in net.features.parameters():
param.requires_grad = False-》冻结特征提取部分所有的权重。

6.features代表的是特征层。

import torch
import torch.nn as nn
from torchvision import transforms, datasets
import json
import os
import torch.optim as optim
from model import MobileNetV2


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

data_transform = {
    "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                 transforms.RandomHorizontalFlip(),
                                 transforms.ToTensor(),
                                 transforms.Normalize([0.485, 0.456