用 Pytorch 训练一个自己的分类器
原文网址：https://zhuanlan.zhihu.com/p/74722043/
环境搭建
首先需要搭建软硬件环境，如果有 GPU 的话那最好，没有的话也没关系，跑 demo 还是可以的。如果数据集大的话还是需要 GPU 做支持，GPU 的训练速度是 CPU 的 10 倍以上。操作系统推荐 Linux，我由于工作需要已经把之前的 Linux 换成了 Windows，就主要介绍 Windows。环境搭建的大致步骤如下，如果碰到问题欢迎在下方留言讨论。
- 安装 python，推荐 python3，本人安装的是 3.7，直接去官网下载 exe 安装即可，要注意的是安装过程中需要勾选 “将其添加到环境变量” 选项，这样就可以直接在命令行输入 python 进入 python 提示符界面了。
- 如果有 GPU 的话需要安装 GPU 对应的驱动以及 CUDA，驱动直接官网找到对应显卡版本下载安装，CUDA 的话直接搜索 CUDA 点击进入系统选择页面选择自己的系统版本 Download，下载完成安装一下就好了。安装完成之后可以在 "C:\Program Files\NVIDIA Corporation\NVSMI" 路径下面运行一下 nvidia-smi.exe 确认安装成功。
- 安装 pytorch，torchvision。在 pytorch 官网主页就可以选择需要的版本以及安装方式，推荐直接 pip 安装，两行命令搞定。
安装完成之后在命令行里验证一下有没有实际安装成功，成功的话应该跟我一样：
训练模型
训练模型中最重要的就是训练集的准备，模型就像是一个小孩子，一开始他啥也不知道，训练的过程就是在“教”他一些。要是一开始“教”的就是错的，那么也不可能期望他能在考试的时候把题目答对是不是。训练集的准备通常需要耗费大量人力物力，所以现在正在往半监督或无监督的方向发展，这是后话。啰嗦这么多，其实我就是想强调训练集的重要性，因为之前吃过亏，在这里提醒一下大家。
这里我使用开放数据集做为例子:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=16, shuffle=True, num_workers=1)
valset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
valloader = torch.utils.data.DataLoader(valset, batch_size=16, shuffle=False, num_workers=1)
datasets = {"train":trainset, "val":valset}
dataloaders = {"train":trainloader, "val":valloader}torchvision 中集成了一些开放数据集，可以直接下载。上面的代码创建了训练集和验证集的数据加载器，batch_size 表示每个 batch 中图片的数量，如果显存大的话可以设置大一点如(32/64/128)，shuffle 表示是否打乱数据集，在训练的时候需要打乱，验证的时候自然不需要，num_workers 表示加载数据集的进程数，需要注意在 Windows 上只能设置为1，否则会报错。在 Linux 上可以设置得大一点加快训练速度。
你也可以定义自己的数据集，只需要继承torch.utils.data.Dataset，然后实现一下自己的 __getitem__() 和 __len__() 就可以。下面是一个最简单的例子，你可以根据自己的需求定制:
class MyDataset(Dataset):
def __init__(self, image_path, transform=None):
self.image_path = image_path
self.transform = transform
def __len__(self):
return 1
def __getitem__(self, index):
pic = Image.open(self.image_path).convert("RGB")
if self.transform:
pic = self.transform(pic)
return (index, pic)训练集准备结束，可以开始编写训练代码:
model = torchvision.models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 10)模型我们选用 torchvision 中集成的预训练好的 Resnet-18 模型，想要了解更多有关 Resnet 可以看看我的另一篇 https://zhuanlan.zhihu.com/p/74230238 ，因为这个数据集的输出有 10 种类型，所以最后全连接层的输出改成了 10。
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)定义损失函数及优化器，这里采用了 交叉熵 作为最优化的目标，优化器采用 SGD，初始学习率为 0.01，动量 0.9，这些都是比较常用的参数值。
# GPU 是否可用
cuda = torch.cuda.is_available()
if cuda:
model.cuda()
best_accuracy = 0.0
start_time = time.time()
epoches = 5
for epoch in range(epoches):
print('Epoch {}/{}'.format(epoch, epoches - 1))
print('-' * 40)
since_epoch = time.time()
for phase in ["train", "val"]:
if phase == "train":
model.train()
else:
model.eval()
running_loss = 0.0
running_corrects = 0
for data in dataloaders[phase]:
inputs, labels = data
# put data on GPU
if cuda:
inputs = inputs.cuda()
labels = labels.cuda()
# init optimizer
optimizer.zero_grad()
# forward
outputs = model(inputs)
_, preds = torch.max(outputs.data, 1)
# loss
loss = criterion(outputs, labels)
if phase == "train":
# backward
loss.backward()
# update params
optimizer.step()
# total loss
running_loss += loss.item() * inputs.size(0)
# correct numbers
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / len(datasets[phase])
epoch_acc = float(running_corrects) / len(datasets[phase])
time_elapsed_epoch = time.time() - since_epoch
print('{} Loss: {:.4f} Acc: {:.4f} in {:.0f}m {:.0f}s'.format(
phase, epoch_loss, epoch_acc, time_elapsed_epoch // 60