【深度学习实战】Pytorch 修改内置 Resnet 实现 Caltech256 数据集分类训练

1 背景

之前介绍过如何自定义一个神经网络来训练数据。通常,对于一般工作者来说,都是选择一些成名的网络,比如VGG, Resnet, Googlenet等。通过改造这些已经被证实过好用的网络来实现我们自己的功能。这里我们以Resnet50为例。

Caltech256数据集:是加利福尼亚理工学院收集整理的数据集,该数据集选自Google Image 数据集,并手工去除了不符合其类别的图片。在该数据集中,图片被分为256类,每个类别的图片超过80张。这里我将Caltech256数据集按同分布划分成了训练集和验证集,已上传到kaggle,大家可以到kaggle下载

2 准备数据集。

加载工具包

import torch from torch import nn from torch.utils.data import Dataset, DataLoader import torch.nn.functional as F import torchvision from torchvision import transforms from PIL import Image from tqdm import tqdm import os import glob

2.1 自定义 Dataset

  • 使用 glob 来扫描数据集下的所有图片
  • 输入 类名列表来保证训练集和验证集的类名所对应的索引是一致的
class Caltech256(Dataset): def __init__(self, img_dir, classes, transform = None): ''' Params: img_dir: str, 数据集所在路径 classes: list, 数据集包含的所有类名 transform: transform, 数据处理方法 Attributes: img_pths: list, 数据集所包含的所有图片路径的集合 classes: list, 数据集包含的所有类名 transform: transform, 数据处理方法 ''' img_pths = sorted(glob.glob(img_dir + os.sep + '**' + os.sep + '**.jpg')) # 34745 张图片 assert img_pths, 'no jpg file in ' + img_dir self.img_pths = img_pths self.classes = classes self.transform = transform def __len__(self): return len(self.img_pths) def __getitem__(self, idx): img_pth = self.img_pths[idx] cls_name = img_pth.split(os.sep)[-2] label = self.classes.index(cls_name) image = Image.open(img_pth).convert("RGB") if self.transform: image = self.transform(image) return image, label

2.2 数据增强和加载

  • 由于 Caltech256 数据集是将每个类的图片放在同一文件夹下。所以我们使用os.listdir()来获取所有类名,然后进行排序重建索引以达到每次类目所对应的索引都是一致的目的。
  • 数据增强:我们选择一些简单且不会影响图片标签的数据增强方法
    • 水平翻转
    • 垂直翻转
    • 随机剪裁,为了避免剪裁到的图片大部分为背景,所以我们先resize成 300x300 的图片再进行裁剪
    • 随机擦除
  • num_workers: 在windows系统下需设置为 0 不然程序无法运行
train_dir = './data/caltech256/train' val_dir = './data/caltech256/val' classes = sorted(os.listdir(train_dir)) # class names transform = transforms.Compose([ transforms.Resize(size=300), # 先resize成一个 300x300 的图片 transforms.RandomHorizontalFlip(p=0.5), # 水平翻转 transforms.RandomVerticalFlip(p=0.3), # 垂直翻转 transforms.ToTensor(), # 转化为 Tensor 并归一化 transforms.RandomApply(torch.nn.ModuleList([transforms.RandomCrop(size=227)]), p=0.5),# 随机裁剪,发生的概率为0.5 transforms.Resize(size=(227, 227)), transforms.RandomErasing(p=0.4) ]) train_set = Caltech256(train_dir, classes, transform = transform) val_set = Caltech256(val_dir, classes, transform = transform) train_loader = DataLoader(train_set, batch_size=64, num_workers=0, shuffle=True) val_loader = DataLoader(val_set, batch_size=64, num_workers=0, shuffle=True)

3 准备模型

3.1 加载 Pytorch 预训练模型, 并查看模型信息(中间部分被省略)

model = torchvision.models.resnet50(pretrained=True) print(model)
ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  ...
  ...
  ...
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=2048, out_features=1000, bias=True)
)

3.2 修改模型

我们可以看到模型的输出层是输出 1000 个,我们把它改成 256 以符合我们的数据集

fc_features = model.fc.in_features # 获取fc层的输入维度 model.fc = nn.Linear(fc_features, 256) # 将分类层改成 256 个输出 print(model) # 查看模型信息
ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  ...
  ...
  ...
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=2048, out_features=256, bias=True)
)

4 编写训练代码

训练代码一共有两个方法,一个是训练数据,另一个是验证数据。我们使用 tqdm 来显示我们的训练进度

  • 训练数据
def train_loop(dataloader, model, loss_fn, optimizer, epoch, epochs): pbar = tqdm(dataloader) # 使用进度条 print(('\n' + '%10s' * 3) % ('Epoch', 'gpu_mem', 'loss')) for X, y in pbar: # gpu 加速 if torch.cuda.is_available(): X = X.cuda() y = y.cuda() # 正向传播,并计算loss pred = model(X) loss = loss_fn(pred, y) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() loss = loss.item() # 查看 gpu 内存 mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # 打印日志 s = ('%10s' * 3) % ('%g/%g' % (epoch + 1, epochs), mem, '%.7g'% loss) pbar.set_description(s)
  • 验证数据
def test_loop(dataloader, model, loss_fn, epoch, epochs): size = len(dataloader.dataset) correct = 0 pbar = tqdm(dataloader) with torch.no_grad(): for X, y in pbar: # 使用gpu加速 if torch.cuda.is_available(): X = X.cuda() y = y.cuda() pred = model(X) loss = loss_fn(pred, y).item() correct += (pred.argmax(1) == y).type(torch.float).sum().item() mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) s = ('%10s' * 3) % ('%g/%g' % (epoch + 1, epochs), mem, '%.7g'% loss) pbar.set_description(s) correct /= size print(f'Accuracy: {(100*correct):>0.1f}% \n') return correct

5. 进行训练

通常一个深度学习训练包含以下四个部分

  • 数据集:Caltech256
  • 模型: Resnet50
  • 优化器: Adam
  • 损失函数: 交叉熵损失函数(CrossEntropyLoss)
use_gpu = torch.cuda.is_available() # 检查 gpu 是否可用 model = torchvision.models.resnet50(pretrained=True) fc_features = model.fc.in_features model.fc = nn.Linear(fc_features, 256) loss_fn = nn.CrossEntropyLoss() # 损失函数 if use_gpu: model = model.cuda() loss_fn = loss_fn.cuda() optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999)) # 优化器 # 训练 100 个epochs epochs = 100 best = 0.0 for t in range(epochs): train_loop(train_loader, model, loss_fn, optimizer, t, epochs) correct = test_loop(val_loader, model, loss_fn, t, epochs) if correct > best: best = correct torch.save(model, './best.pth') # 保存在验证集准确率最高的模型 print("Done!")

大约进行30-40个epochs可以得到 65% 的准确率

阅读(117)
评论(0)
updated@2021-06-17
评论区
目录