pytorch基础教学简单实例（附代码）

目录

1.简介 

2.数据集

3.模型初始化

4.训练参数

5.训练&验证

6.保存&加载模型

---

1.简介 

这篇文章主要是针对刚入门pytorch的小伙伴，会带大家完整走一遍使用神经网络训练的流程，以及介绍一些pytorch常用的函数。如果还未安装pytorch或者安装有困难，可以参考我的上一篇文章：

Windows Anaconda精简安装cuda+pytorch+torchvision

---

2.数据集

这里使用的是FashionMNIST，因为可以直接调用pytorch里的代码下载比较方便。加载代码如下：

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# ----------数据集----------

# 加载MNIST数据集的训练集
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)
# 加载MNIST数据集的测试集
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

# batch大小
batch_size = 64

# 创建dataloader
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

# 遍历dataloader
for X, y in test_dataloader:
    print("Shape of X [N, C, H, W]: ", X.shape)     # 每个batch数据的形状
    print("Shape of y: ", y.shape)                  # 每个batch标签的形状
    break

---

3.模型初始化

手动搭建了两个全连接层的神经网络结构。代码如下：

# ----------模型----------
# 定义模型
class NeuralNetwork(nn.Module):
    def __init__(self):     # 初始化，实例化模型的时候就会调用
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()     # [64, 1, 28, 28] -> [64, 1*28*28]
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),      # [64, 1*28*28] -> [64, 512]
            nn.ReLU(),
            nn.Linear(512, 512),        # [64, 512] -> [64, 512]
            nn.ReLU(),
            nn.Linear(512, 10)          # [64, 512] -> [64, 10]
        )

    def forward(self, x):   # 前向传播，输入数据进网络的时候才会调用
        x = self.flatten(x)                     # [64, 1*28*28]
        logits = self.linear_relu_stack(x)      # [64, 10]
        return logits

# 使用gpu或者cpu进行训练
device = "cuda" if torch.cuda.is_available() else "cpu"
# 打印使用的是gpu/cpu
print("Using {} device".format(device))
# 实例化模型
model = NeuralNetwork().to(device)
# 打印模型结构
print(model)

---

4.训练参数

损失函数用来计算预测输出和真实输出的差值，优化器用来更新网络的权重参数，scheduler用来调整训练过程中的学习率。代码如下：

# ----------训练参数设置----------
loss_fn = nn.CrossEntropyLoss()                             # 损失函数设置
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)    # 学习率设置
epochs = 5                                                  # 训练迭代次数设置

---

5.训练&验证

训练和测试的步骤可以总结为下图，具体代码如下：

# 训练函数
def train(train_dataloader, model, loss_fn, optimizer):
    """
        训练网络

        输入:
            train_dataloader:   训练集的dataloader
            model:              网络模型
            loss_fn:            损失函数
            optimizer:          优化器
        """
    # 切换到train模式
    model.train()
    # 遍历dataloader
    for images, labels in train_dataloader:
        # 将数据和标签加载到device上
        images, labels = images.to(device), labels.to(device)
        # 输入数据到模型里得到输出
        pred = model(images)
        # 计算输出和标签的loss
        loss = loss_fn(pred, labels)
        # 反向推导
        optimizer.zero_grad()
        loss.backward()
        # 步进优化器
        optimizer.step()


# 测试函数
def test(test_dataloader, model, loss_fn):
    """
        测试网络

        输入:
            test_dataloader:    测试集的dataloader
            model:              网络模型
            loss_fn:            损失函数
        """
    # 测试集大小
    size = len(test_dataloader.dataset)
    # 测试集的batch数量
    num_batches = len(test_dataloader)
    # 切换到测试模型
    model.eval()
    # 记录loss和准确率
    test_loss, correct = 0, 0
    # 梯度截断
    with torch.no_grad():
        for images, labels in test_dataloader:  # 遍历batch
            # 加载到device
            images, labels = images.to(device), labels.to(device)
            # 输入数据到模型里得到输出
            pred = model(images)
            # 累加loss
            test_loss += loss_fn(pred, labels).item()
            # 累加正确率
            correct += (pred.argmax(1) == labels).sum().item()
    # 计算平均loss和准确率
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

6.保存&加载模型

在训练好网络后，我们可以将训练好的权重保存下来，这样下次要使用的时候就不用再次训练，直接加载就好，代码如下：

​

# 保存模型
torch.save(model.state_dict(), "model.pth")
# 加载模型
model = NeuralNetwork()
model.load_state_dict(torch.load("model.pth"))

---

完整的代码可以在我的github上找到：lizhiTech/pytorch_simple_teaching_example: 零基础pytorch深度学习安装使用代码实例新手教学 (github.com)

我们也提供包括深度学习、计算机视觉、机器学习等其他方向的其他代码及辅导服务，有需求可以联系我们，详情见下图：

​