使用PyTorch搭建神经网络

使用PyTorch搭建神经网络

import torch
import torch.nn as nn
from torchvision import transforms, datasets
from torch.utils.data import DataLoader
import time

train_start = time.time()  # 训练开始时间
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

网络构建

class SimpleNN(nn.Module):  
    def __init__(self):  
        super(SimpleNN, self).__init__()  
        # 定义第一层全连接层：输入特征784，输出特征128  
        self.fc1 = nn.Linear(784, 128)  
        # 定义第二层全连接层：输入特征128，输出特征128  
        self.fc2 = nn.Linear(128, 128)  
        # 定义输出层：输入特征128，输出特征10（对应10个数字类别）  
        self.output = nn.Linear(128, 10)  
  
    def forward(self, x):  
        # 对第一层的输出应用ReLU激活函数  
        x = torch.relu(self.fc1(x))  
        # 对第二层的输出也应用ReLU激活函数  
        x = torch.relu(self.fc2(x))  
        # 通过输出层得到最终的分类结果  
        x = self.output(x)  
        return x  # 确保 return 语句正确缩进

数据加载与处理

# 数据加载与处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图片数据转换为Tensor
    transforms.Normalize((0.5,), (0.5,)),  # 对数据进行归一化处理
]) 

加载训练数据集

train_set = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)

加载测试数据集

test_set = datasets.MNIST(root='./data', train=False, transform=transform)
test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

将模型实例化并转移到定义的设备

model = SimpleNN().to(device)

选择优化器和损失函数

optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

训练过程

for epoch in range(5):  # 训练5个循环周期  
    model.train()  # 设置为训练模式（对于某些模块如 Dropout 和 BatchNorm 是必要的）  
    for images, labels in train_loader:  
        # 调整图片形状并转移到相同的设备  
        images, labels = images.view(-1, 28*28).to(device), labels.to(device)  
          
        # 清除历史梯度  
        optimizer.zero_grad()  
          
        # 前向传播计算模型输出  
        output = model(images)  
          
        # 计算损失  
        loss = criterion(output, labels)  
          
        # 反向传播计算梯度  
        loss.backward()  
          
        # 更新模型参数  
        optimizer.step()  
          
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

评估函数

def evaluate_model(model, data_loader, device):  
    model.eval()  # 将模型设置为评估模式  
    total_correct = 0  
    total = 0  
    with torch.no_grad():  # 禁止梯度计算  
        for images, labels in data_loader:  
            # 将图片和标签数据转移到相同的设备  
            images, labels = images.view(-1, 28*28).to(device), labels.to(device)  
            output = model(images)  # 前向传播得到预测结果  
            _, predicted = torch.max(output.data, 1)  # 得到预测的类别  
            total += labels.size(0)  
            total_correct += (predicted == labels).sum().item()  # 统计正确预测的数量  
    accuracy = 100 * total_correct / total  
    print(f'Accuracy: {accuracy:.2f}%')  # 在函数内部打印准确率  
    return accuracy  

调用 evaluate_model 并捕获准确率

accuracy = evaluate_model(model, test_loader, device)  

记录训练结束时间并打印训练时间

train_end = time.time()  
print("训练时间: {:.2f}秒".format(train_end - train_start))```

调用 evaluate_model 并捕获准确率

accuracy = evaluate_model(model, test_loader, device)  

记录训练结束时间并打印训练时间

train_end = time.time()  
print("训练时间: {:.2f}秒".format(train_end - train_start))

参考：https://zhuanlan.zhihu.com/p/696017829