torchvision教程

Torchvision 是 PyTorch 中一个非常有用的工具包，它提供了许多用于计算机视觉的工具和数据集。下面是一个简单的 Torchvision 教程，介绍如何使用 Torchvision 进行图像分类。

1. 安装 Torchvision

Torchvision 可以通过 pip 安装，命令如下：

pip install torchvision

2. 加载数据集

Torchvision 提供了多个常用的数据集，例如 MNIST、CIFAR10、ImageNet 等。这里以 MNIST 数据集为例：

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

train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=True)

上述代码中，我们首先导入了 `datasets` 和 `transforms` 模块，然后使用 `MNIST` 函数加载数据集。`root` 参数指定数据集存储路径，`train` 参数指定是否加载训练集，`transform` 参数指定数据预处理操作（这里使用 `ToTensor` 将图像转换为张量），`download` 参数指定是否自动下载数据集。

3. 定义模型

在这个例子中，我们使用一个简单的卷积神经网络进行图像分类。代码如下：

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(1, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    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 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

上述代码中，我们定义了一个名为 `Net` 的类，它继承自 `nn.Module`，并实现了一个简单的卷积神经网络。在 `__init__` 方法中，我们定义了卷积层、池化层和全连接层，以及它们的输入输出大小。在 `forward` 方法中，我们定义了网络的正向传播过程。

4. 训练模型

训练模型的代码如下：

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 % 2000 == 1999:
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

上述代码中，我们使用交叉熵损失函数和随机梯度下降优化器训练模型。在每个 epoch 中，我们遍历训练集，计算模型输出和损失，然后反向传播更新模型参数。

5. 测试模型

测试模型的代码如下：

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

上述代码中，我们使用测试集测试模型，并计算模型在测试集上的准确率。

以上就是一个简单的 Torchvision 教程，介绍了如何使用 Torchvision 进行图像分类。