神经网络模型代码pytorch

### 使用 PyTorch 编写神经网络模型

#### 创建一个简单的全连接神经网络

为了展示如何使用 `PyTorch` 构建和训练一个简单的多层感知机（MLP），下面提供了一个完整的代码示例。此例子展示了如何定义模型结构、设置损失函数、选择优化算法以及执行前向传播和反向传播。

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


class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        # 定义三层线性变换，即输入层、隐藏层和输出层
        self.fc1 = nn.Linear(784, 500)  
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        # 将二维图像展平成一维张量
        x = x.view(-1, 784)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x


def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()  # 设置模型为训练模式
    running_loss = 0.0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        optimizer.zero_grad()  # 清除之前的梯度
        
        output = model(data)  # 前向传递
        loss = criterion(output, target)  # 计算损失值
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新参数
        
        running_loss += loss.item()

        if batch_idx % 100 == 99:  # 打印每100批次的日志信息
            print(f'Epoch [{epoch}], Batch[{batch_idx + 1}], Loss: {running_loss / 100:.3f}')
            running_loss = 0.0


if __name__ == '__main__':
    epochs = 5
    learning_rate = 0.01
    batch_size = 64
    transform = ToTensor()

    dataset_train = MNIST(root='./data', train=True, download=True, transform=transform)
    dataloader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    net = SimpleNN().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=learning_rate)

    for e in range(epochs):
        train(net, device, dataloader_train, optimizer, criterion, e + 1)

这段代码实现了基于MNIST手写数字识别任务的一个小型卷积神经网络的搭建与训练流程[^2]。该实例中包含了数据预处理、模型设计、损失函数的选择、优化器配置及其迭代更新机制等内容。