pytorch + visdom CNN处理自建图片数据集的方法

对于使用PyTorch训练CNN的过程，一般情况下需要进行图片的预处理、数据集的加载，以及训练过程的可视化等步骤。其中，使用visdom进行训练过程的可视化非常方便，其支持的图形工具非常丰富。

下面，我们将围绕着“pytorch + visdom CNN处理自建图片数据集的方法”，从以下几个方面进行详细讲解。

1.数据集的准备

对于训练CNN所需的数据集，一般情况下需要进行以下几个步骤的准备：

（1）将数据集分为训练集、测试集，一般比例为8:2或7:3；

（2）将数据集中的图片进行缩放，使其大小统一，便于CNN模型的训练；

（3）将数据集中的图片进行标记，对每张图片给出对应的标签。

在PyTorch中，我们可以使用 torchvision.datasets.ImageFolder 对自定义数据集进行加载，它会自动为数据集中的每个样本生成标签，并将其存储在分类文件夹中。示例代码如下：

import torchvision.transforms as transforms
import torchvision.datasets as datasets

# 数据增强
data_transforms = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor(),
    ])

# 数据集加载
image_datasets = datasets.ImageFolder(root='./data', transform=data_transforms)

# 数据集分割
train_size = int(0.8 * len(image_datasets))
test_size = len(image_datasets) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(image_datasets, [train_size, test_size])

2.CNN模型的定义

在PyTorch中定义CNN模型，可以直接继承 nn. Module 类，然后在子类中实现 init 和 forward 方法。其中，init 方法一般用来定义所需的层， forward 方法则用来实现前向传播。示例代码如下：

import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 29 * 29, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 29 * 29)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

3.训练过程的实现

在PyTorch中，我们可以使用 torch.optim 模块中的优化器对CNN模型进行训练。其中，常用的优化器包括 Adam(), SGD() 和 RMSprop() 等。示例代码如下：

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0

4.训练过程的可视化

使用visdom库可以直观的显示训练过程中的Loss值、准确率等变量。示例如下：

import visdom

def plot_loss_acc(loss, acc):
    x = torch.arange(len(loss))
    y1 = loss
    y2 = acc
    assert len(x) == len(y1) == len(y2)
    y1_opts = dict(title='loss', xlabel='epoch', ylabel='Loss')
    y2_opts = dict(title='acc', xlabel='epoch', ylabel='Acc.')
    win = 'result'
    if not vis.win_exists(win):
        vis.line(X=x,Y=y1, win=win,opts=y1_opts, name='train-loss')
        vis.line(X=x,Y=y2, win=win,opts=y2_opts, name='train-accuracy')
    else:
        vis.line(X=x,Y=y1, win=win,opts=y1_opts, update='append', name='train-loss')
        vis.line(X=x,Y=y2, win=win,opts=y2_opts, update='append', name='train-accuracy')

整合以上四个步骤后，即可方便地使用PyTorch和visdom对自定数据集进行CNN模型训练、训练过程的可视化。