基于Pytorch TorchHub和RESNET的图像分类案例

基于Pytorch TorchHub和RESNET的图像分类案例

此章节中通过一个具体案例详细介绍如何使用TorchHub，基于已经训练好的ResNet模型进行迁移学习分类任务。我们将学习这些模型背后的核心思想，并根据我们选择的任务对其进行微调。

Torch Hub在网络上提供了大量经过预先训练的模型权重，可以识别可能出现的所有问题，并通过将整个过程浓缩到一行来解决这些问题。因此，不仅可以在本地系统中加载SOTA模型，还可以选择是否需要对其进行预训练。

• 运行平台：Win10
• IDE：VSCode
• 资源下载：所有代码都包含了详细注释，下载链接：

---

---

1 简介

1.1 回顾 RESNET

传统的CNN网络问题：网络越叠越深之后难以收敛，梯度消失/爆炸在一开始就阻碍网络的收敛。通过标准初始化和中间标准化层在很大程度上解决。这使得数十层的网络能通过具有反向传播的随机梯度下降（SGD）开始收敛。当更深的网络能够开始收敛时，暴露了一个退化问题：随着网络深度的增加，准确率达到饱和然后迅速下降。意外的是，这种下降不是由过拟合引起的，并且在适当的深度模型上添加更多的层会导致更高的训练误差。

残差网络提出residual层，明确地让这些层拟合残差映射，而不是希望每几个堆叠的层直接拟合期望的基础映射。在传统的CNN网络中，其实只存在下图类似F(x)的非线性堆叠层。残差网络将原始的映射重写为F(x)+x。道理也很简单：在极端情况下，如果一个恒等映射是最优的，那么将残差置为零比通过一堆非线性层来拟合恒等映射更容易。快捷连接简单地执行恒等映射，并将其输出添加到堆叠层的输出。恒等快捷连接既不增加额外的参数也不增加计算复杂度。整个网络仍然可以由带有反向传播的SGD进行端到端的训练。

RESNET系列网络结构
（i）对于相同的输出特征图尺寸，层具有相同数量的滤波器；
（ii）如果特征图尺寸减半，则滤波器数量加倍，以便保持每层的时间复杂度。我们通过步长为2的卷积层直接执行下采样。

1.2 文件结构

├── classifier.py       # 包含项目的模型体系结构
├── config.py           # 参数配置文件
├── model               # 已训练好的Resnet18模型文件（存放于此文件夹中）
├── output              # 经过迁移学习的模型，以及plots（存放于此文件夹中）
├── dataset             # 数据集。
│   ├── test_set
│       ├── cats
│       └── dogs
│   └── training_set
│       ├── cats
│       └── dogs
├── main_train.py       # 模型训练相关主代码
├── main_inference.py   # 模型inference相关主代码
└── requirements.txt    # 依赖文件的安装

1.3 搭建虚拟环境

我们可以自行选择使用Anaconda或者virtualenv来新建虚拟环境。

1.3.1 Anaconda

首先，我们在Windows的平台下安装Anaconda3。具体的安装步骤此处略过，参见Anaconda的官方文档。

安装完后，新建虚拟环境。使用conda create -n your_env_name python=X.X（3.6、3.8等)命令创建python版本为X.X、名字为your_env_name的虚拟环境。

这里我输入了conda create -n torchhubex python=3.8。

安装完默认的依赖后，我们进入虚拟环境：conda activate torchhubex。注意，如果需要退出，则输入conda deactivate。另外，如果Terminal没有成功切换到虚拟环境，可以尝试conda init powershell，然后重启terminal。

然后，我们在虚拟环境中下载好相关依赖：pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple。需要安装的依赖如下：

matplotlib==3.5.1
scikit-learn==1.0.2
opencv-contrib-python==4.5.5.64
torch==1.11.0
torchvision==0.12.0
torchaudio==0.11.0
tqdm==4.63.0

1.3.2 virtualenv

在terminal中安装virtualenv：pip3 install virtualenv -i https://pypi.tuna.tsinghua.edu.cn/simple。然后，我们在选定的路径下安装虚拟环境：virtualenv [venv]。[venv]指的是虚拟环境的名称。完成此指令后，我们就能在此路径下看到一个名为[venv]的文件夹。

进入虚拟环境：.\[venv]\Scripts\activate；退出虚拟环境：deactivate；安装所有的依赖：pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple。

2 代码解读

2.1 数据集使用

此项目使用的数据集源于Kaggle Dogs & Cat。不方便下载的同学我们这边也提供百度网盘链接：https://pan.baidu.com/s/11JMbou33R3fPJ5z0bTNu8A，密码：7hku。

下载后，请解压并将training_set以及test_set放入dataset文件夹下。这个数据集比较简单，实际上就是分类猫和狗两类。

需要运行的主代码为main.py，接下来我们主要解读此文件。

2.2 模型，数据集导入与预处理

torch hub 有一个很简便的地方就在于，其一句话就可以导入训练好的模型：

baseModel = torch.hub.load(config.TORCHHUB_MODELDIR, config.TORCHHUB_MODELNAME, pretrained=True, skip_validation=True)

这里所有参数都在config文件中设置，config.TORCHHUB_MODELDIR的值为pytorch/vision:v0.10.0，config.TORCHHUB_MODELNAME的值为resnet18。这里需要注意，运行这个指令可能会超时，这里我们也已经把这个模型下载下来了，存放于model文件夹下（Win10环境下，将这两个文件放于C:\Users\XXX\.cache\torch\hub目录下，然后再运行baseModel = torch.hub.load(config.TORCHHUB_MODELDIR, config.TORCHHUB_MODELNAME, pretrained=True, skip_validation=True)，应该可以跳过下载环节直接运行）。

接下来，我们需要冻结除了最后一层的其他层参数。迁移学习只学习最后一层。我们的做法是遍历RESNET18的所有层（一共八层，这里没有softmax层），除了最后一层，其他层的param.requires_grad值都设置为False。

currentLayer = 1
for child in baseModel.children():
  if currentLayer < config.RESNET18TOTALLAYERS:  # 如果不是最后一层，那么其参数就冻结
    for param in child.parameters():
      param.requires_grad = False
  else:                                   # 最后一层，迁移学习就学习这层参数
    break
  currentLayer += 1   

接下来就是数据预处理，PyTorch给我们提供了相当简便的模块：torchvision.transforms.Compose

trainTransform = Compose([
  RandomResizedCrop(config.IMAGE_SIZE),
  RandomHorizontalFlip(),
  RandomRotation(90),
  ToTensor(),
  Normalize(mean=config.MEAN, std=config.STD)
])

torchvision.transforms是pytorch中的图像预处理包，一般用Compose把多个步骤整合到一起。这个预处理包里包括了：

• Resize：把给定的图片resize到给定的尺寸
• Normalize：归一化，基于tensor image的mean和std值
• ToTensor：convert a PIL image to tensor (HWC) in range [0,255] to a torch.Tensor(CHW) in the range [0.0,1.0]
• RandomHorizontalFlip：以0.5的概率水平翻转给定的PIL图像
• RandomVerticalFlip：以0.5的概率竖直翻转给定的PIL图像
• RandomRotation：随机旋转（一定角度）
• CenterCrop：在图片的中间区域进行裁剪
• ColorJitter: 随机改变图像的亮度对比度和饱和度
• RandomResizedCrop：将PIL图像裁剪成任意大小和纵横比
• RandomCrop：在一个随机的位置进行裁剪
• Grayscale：将图像转换为灰度图像
• RandomGrayscale：将图像以一定的概率转换为灰度图像
• Pad：填充

最后就是数据的导入了。这里，我们将导入的训练集分为训练集和验证集。

trainDataset = ImageFolder(config.TRAIN_PATH, trainTransform)
(trainDataset, valDataset) = train_val_split(dataset=trainDataset)
trainLoader = DataLoader(trainDataset, batch_size=config.BATCH_SIZE, shuffle=True)
valLoader = DataLoader(valDataset, batch_size=config.BATCH_SIZE, shuffle=True)

2.3 Classifier类

我们基于这个类实现迁移学习，即，我们将RESNET模型与softmax全连接层链接起来，对图像进行分类任务。

class Classifier(Module):
  def __init__(self, baseModel, numClasses, model):
    super().__init__()
    self.baseModel = baseModel		# 这里指RESNET主干网络
    self.fc = Linear(baseModel.fc.out_features, numClasses)

  def forward(self, x):
    '''
    将输入x前向传递，经过RESNET主干网络，以及全连接层，得到classification结果
    '''
    features = self.baseModel(x)
    logits = self.fc(features)
    return logits

这个Classifier类继承于torch.nn.Module。

在main.py文件中，我们对Classifier进行实例化：

model = Classifier(baseModel=baseModel.to(config.DEVICE), numClasses=2, model="resnet")
model = model.to(config.DEVICE)

并且，在进行训练（以及验证）之前，我们需要初始化损失函数，optimizer，softmax类的实例化，

# 初始化Loss function
lossFunc = CrossEntropyLoss()
lossFunc.to(config.DEVICE)
# 初始化optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=config.LR)
# softmax类实例化
softmax = Softmax()

# 计算每个epoch需要的步数（假定需要遍历所有数据的话）
trainSteps = len(trainDataset) // config.BATCH_SIZE
valSteps = len(valDataset) // config.BATCH_SIZE

# 初始化精度与损失函数值（每个循环都会最终记录）
H = {
  "trainLoss": [],
  "trainAcc": [],
  "valLoss": [],
  "valAcc": []
}

2.4 训练

训练代码如下：

model.train()		    # set the model in training mode
totalTrainLoss = 0	# 初始化训练的损失值
totalValLoss = 0	  # 初始化验证的损失值
trainCorrect = 0	  # 初始化在训练过程中正确预测的数量
valCorrect = 0		  # 初始化在验证过程中正确预测的数量
# 遍历训练数据集，计算训练损失值，以及正确预测的数量
for (image, target) in tqdm(trainLoader):
  # send the input to the device
  (image, target) = (image.to(config.DEVICE),	target.to(config.DEVICE))
  # 前向传递，计算训练损失值
  logits = model(image)
  loss = lossFunc(logits, target)
  # gradients值归零
  optimizer.zero_grad()
  # Backpropagation 反向传递
  loss.backward()
  # 更新权重值
  optimizer.step()
  # 计算总的训练损失值
  totalTrainLoss += loss.item()
  # softmax获得预测值
  pred = softmax(logits)
  # 计算正确预测的数量
  trainCorrect += (pred.argmax(dim=-1) == target).sum().item()

在模型训练的for循环中，我们首先将模型设置为训练模式（model.train()）。接下来，我们初始化训练损失（totalTrainLoss）、验证损失（totalValLoss）、训练和验证精度变量（trainCorrect，valCorrect）。进入for循环，我们首先将数据和目标加载到设备（这里指的是cpu或者gpu）。接下来，我们只需通过模型正向传播并获得输出（logits = model(image)），然后将预测和目标插入损失函数（loss = lossFunc(logits, target)），然后是标准的反向传播步骤，我们将梯度归零，执行反向传播（loss.backward()），并更新权重（optimizer.step()）。接下来，我们将损失添加到总训练损失中，通过softmax将模型输出传递到单独的预测值（softmax(logits)），然后将其添加到trainCorrect变量中。

2.5 验证

验证的代码和训练几乎一样，代码如下：

# 验证集计算，关闭autograd模式
with torch.no_grad():
  model.eval()		# set the model in evaluation mode
  # 遍历验证集
  for (image, target) in tqdm(valLoader):
    # send the input to the device
    (image, target) = (image.to(config.DEVICE),	target.to(config.DEVICE))
    # 基于验证集进行预测，并计算损失值。
    logits = model(image)
    valLoss = lossFunc(logits, target)
    totalValLoss += valLoss.item()
    # softmax获得预测值
    pred = softmax(logits)
    # 计算正确预测的数量
    valCorrect += (pred.argmax(dim=-1) == target).sum().item()

2.6 出图

我们将每一个epoch的训练损失，平均验证损失，平均训练正确率以及平均验证正确率都记下来：

# 计算平均训练损失以及平均验证损失
avgTrainLoss = totalTrainLoss / trainSteps
avgValLoss = totalValLoss / valSteps
# 计算平均训练正确率以及验证正确率
trainCorrect = trainCorrect / len(trainDataset)
valCorrect = valCorrect / len(valDataset)
H["trainLoss"].append(avgTrainLoss)
H["valLoss"].append(avgValLoss)
H["trainAcc"].append(trainCorrect)
H["valAcc"].append(valCorrect)

最后画出这些数值随着epoch的变化图。

2.6 Inference

当我们完成了模型迁移学习后，我们可以运行main_inference.py，使用测试数据集来验证我们训练后的模型性能。

main_inference.py的逻辑和main_train.py差不多。一开始，我们需要导入测试数据，以及数据预处理。这里需要注意的是，由于我们在训练模型的时候，对输入的训练数据集进行了计算mean和std，那么也就意味着，当我们使用测试数据集的时候，我们需要用相同的mean和std值对测试数据集进行归一化：

# 数据预处理
testTransform = Compose([
  Resize((config.IMAGE_SIZE, config.IMAGE_SIZE)),
  ToTensor(),
  Normalize(mean=config.MEAN, std=config.STD)
])
# 导入测试数据集
testDataset = ImageFolder(config.TEST_PATH, testTransform)
# 初始化测试数据DataLoader
testLoader = DataLoader(testDataset, batch_size=config.PRED_BATCH_SIZE, shuffle=True)

接下来，我们需要导入RESNET的模型（我们需要导入训练完的模型）

# 加载RESNET 18模型。ref: https://pytorch.org/docs/stable/hub.html#torch.hub.load
baseModel = torch.hub.load(config.TORCHHUB_MODELDIR, config.TORCHHUB_MODELNAME, pretrained=True, skip_validation=True)
# 自定义模型
model = Classifier(baseModel=baseModel.to(config.DEVICE), numClasses=2)
model = model.to(config.DEVICE)
# 读取我们训练好的模型
model.load_state_dict(torch.load(config.MODEL_PATH))

之后，我们初始化损失值以及相关参数：

lossFunc = nn.CrossEntropyLoss()
lossFunc.to(config.DEVICE)
testCorrect = 0
totalTestLoss = 0
softmax = Softmax()

然后就开始把模型跑一下测试数据集了。

with torch.no_grad():
  model.eval()
  # 遍历测试集
  for (image, target) in tqdm(testLoader):
    (image, target) = (image.to(config.DEVICE),	target.to(config.DEVICE))
    logit = model(image)
    loss = lossFunc(logit, target)
    totalTestLoss += loss.item()
    pred = softmax(logit)
    testCorrect += (pred.argmax(dim=-1) == target).sum().item()
    
print("测试正确率: ", testCorrect/len(testDataset))

3 结果

3.1 Train 结果

我们在主路径下运行python .\train_resnet.py（记得在VM环境下运行），在terminal中我们看到：

Downloading: "https://github.com/pytorch/vision/archive/v0.10.0.zip" to C:\Users\gugut/.cache\torch\hub\v0.10.0.zip
Downloading: "https://download.pytorch.org/models/resnet18-f37072fd.pth" to C:\Users\gugut/.cache\torch\hub\checkpoints\resnet18-f37072fd.pth
100%|███████████████████████████████████████████████████████| 44.7M/44.7M [00:18<00:00, 2.49MB/s] 
[INFO] training the network...
  0%|                                                                     | 0/50 [00:00

如果我们导入模型使用

baseModel = torch.hub.load("pytorch/vision:v0.10.0", "resnet18", pretrained=True, skip_validation=True)

可能会遇到超时的报错：TimeoutError: [WinError 10060] 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。，因为模型的调用是从国外的server链接上下载下来的。

最后，保存的图片为：

3.2 Inference 结果

当我们首先运行完main_train.py后，通过迁移学习得到模型后，我们再运行main_test.py，在terminal的日志如下：

Using cache found in C:\Users\XXX/.cache\torch\hub\pytorch_vision_v0.10.0
  0%|                                                                     | 0/32 [00:00