Quick Start

在安装后Pytorch之后，官方提供了一个基于FashionMNIST的分类任务，通过此任务可以让我们快速入门怎么样使用数据集去训练一个神经网络模型

数据集加载

torchvision中集成了datasets的API以便于快速的下载一些常用的数据集

1
2


from torchvision import datasets
from torchvision.transforms import ToTensor

使用datasets

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


# Download training data from open datasets.
train_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# Download test data from open datasets.
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

root是存储下载数据集的路径。ToTensor是把原始数据集转换为Pytorch的张量，一般来说，对于一张原始图片，其shape就是[H, W, C]，在Pytorch中，为了方便计算，通常在ToTensor后会把其shape转换为[C, H, W]的形状。

加载数据集

在Pytorch中，加载数据集都是使用DataLoader类来进行加载，返回的是一个类似于迭代器的数据类型，值得注意的是，我们也可以自行构建一个自己的数据集来进行数据加载，而且这个数据集必须是Dataset类的子类。

1

from torch.utils.data import DataLoader

DataLoader的第一个参数类型就是Dataset类

1

train_dataloader = DataLoader(train_data, shuffle= True, batch_size= 64)

后面就是一些必选参数，像shuffle, batch_size这些需要自己选择，一般来说，在训练集需要打开shuffle来进行训练，而测试集不需要随机打乱。batch_size就是每一次选择几张图片作为一个batch来进行训练，batch_size选择的越大，占用的内存也就越高。

1

test_dataloader = DataLoader(test_data, shuffle=True, batch_size= 64)

加载之后，输入的数据应该有四个维度，即[B, C, H, W]，对于FashionMNIST来说，其数据维度就是[64, 1, 28, 28]，我们可以来验证一下

1
2
3
4


for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

选择损失函数和优化器

对于多分类问题，可以选择交叉熵损失函数即

1

loss_func = nn.CrossEntropyLoss()

同样的，我们也可以选择不同的优化方法，例如SGD，SGD+Momentum，Adam等等这些损失器

1

optim = torch.optim.Adam(model.parameters(), lr=0.001)

模型的训练

开始训练模型，可以把模型的训练封装为一个函数

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24


def train(dataloader, model, loss_func, optimizer):
    # 获取总的数据个数
    size = len(dataloader.dataset)
    # 开启训练模式
    model.train()
    for batch, (x, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        # 通过model来获取pred_y
        pred = model(X)
        # 损失函数
        loss = loss_func(pred, y)
        
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
        # 梯度归零
        optimizer.zero_grad()
        # 每100个batch得到一个输出
        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f'loss: {loss:>7f} [{current:>5d} / {size:>5d}]')
        
    

测试阶段

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


def test(dataloader, model, loss_fn):
    # 获取总的样本个数
    size = len(dataloader.dataset)
    # 获取批次数
    num_batches = len(dataloader)
    # 开启评估模式
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            # 获取验证集上的损失
            test_loss += loss_fn(pred, y).item()
            # 正确的样本率
            # argmax(1) 与 真值是否相等
            correct += (pred.argmax(dim=1) == y).sum().item()
        # 获取测试集上的损失值
        test_loss /= num_batches