PyTorch 深度学习项目结构及nn.Module介绍

PyTorch 深度学习项目结构介绍及nn.Module介绍

在 PyTorch 深度学习项目中，通常会涉及到模型定义（Network）、数据处理（Dataset）、优化器（Optimizer）、训练过程（Run）和测试过程（Test），这些部分的清晰划分有助于项目的扩展和维护。

1. 项目结构

my_pytorch_project/
│
├── data/                        # 存放数据集的目录
│   ├── train/                   # 训练数据
│   ├── val/                     # 验证数据
│   └── test/                    # 测试数据
│
├── models/                      # 存放模型定义的目录
│   └── network.py               # 网络结构（模型定义）
│
├── transforms/                  # 存放数据预处理相关的变换
│   └── transform.py             # 数据预处理和增强（如归一化、翻转等）
│
├── utils/                       # 存放辅助工具函数的目录
│   ├── dataset.py               # 数据集处理（继承自`Dataset`类）
│   ├── optimizer.py             # 优化器设置
│   ├── run.py                   # 训练和验证过程
│   ├── test.py                  # 测试过程
│   └── utils.py                 # 其他工具函数（如保存模型、加载模型等）
│
├── output/                      # 存放训练输出的目录（模型权重、日志等）
│   ├── checkpoints/             # 保存模型检查点
│   └── logs/                    # 存放训练日志文件
│
├── requirements.txt             # 项目依赖的库
├── train.py                     # 训练脚本
├── evaluate.py                  # 测试脚本
└── README.md                    # 项目说明文档

2. 代码文件详细介绍

2.1 network.py（模型定义）

模型（Network）通常是 PyTorch 中通过 nn.Module 继承定义的。在这个例子中，我们使用了一个卷积神经网络（CNN）。

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 卷积层
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # 假设输入图片大小为28x28
        self.fc2 = nn.Linear(128, 10)  # 输出10类分类
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)  # 最大池化
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = x.view(-1, 64 * 7 * 7)  # 扁平化
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

2.2 transform.py（数据预处理）

Transform 模块包含了图像预处理和增强的逻辑，常见的包括归一化、翻转、裁剪等。

from torchvision import transforms

def get_transforms():
    return transforms.Compose([
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.RandomRotation(10),      # 随机旋转
        transforms.ToTensor(),              # 转换为Tensor
        transforms.Normalize((0.5,), (0.5,)) # 归一化
    ])

2.3 dataset.py（数据集处理）

Dataset 是 PyTorch 中的核心类，用于处理自定义数据集。在这个文件中，我们可以定义如何读取和处理数据。

import torch
from torch.utils.data import Dataset
from torchvision import datasets, transforms
from transforms.transform import get_transforms

class MyDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data = datasets.MNIST(data_dir, train=True, download=True, transform=transform)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        img, label = self.data[idx]
        return img, label

2.4 optimizer.py（优化器设置）

在 optimizer.py 中定义了优化器（如 Adam、SGD 等）的设置和调整。

import torch.optim as optim

def get_optimizer(model):
    return optim.Adam(model.parameters(), lr=0.001)

2.5 run.py（训练和验证）

Run 脚本负责训练和验证过程，计算损失并调整权重。

import torch
from torch.utils.data import DataLoader
from utils.dataset import MyDataset
from models.network import SimpleCNN
from utils.optimizer import get_optimizer
from utils.utils import save_checkpoint
from transforms.transform import get_transforms

def train(model, train_loader, criterion, optimizer, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        running_loss = 0.0
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()

        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader)}')
        save_checkpoint(model, epoch, f'output/checkpoints/model_epoch_{epoch+1}.pth')

def run():
    # 数据加载
    transform = get_transforms()
    train_dataset = MyDataset('data/train', transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

    # 模型定义
    model = SimpleCNN()

    # 损失函数和优化器
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = get_optimizer(model)

    # 训练过程
    train(model, train_loader, criterion, optimizer, num_epochs=10)

2.6 test.py（测试过程）

Test 脚本用于在训练后的模型上评估其性能，通常包括准确度计算。

import torch
from models.network import SimpleCNN
from utils.dataset import MyDataset
from torch.utils.data import DataLoader
from utils.utils import load_checkpoint

def test(model, test_loader):
    model.eval()  # 设置为评估模式
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Accuracy: {accuracy:.2f}%')

def run_test():
    # 加载模型
    model = SimpleCNN()
    checkpoint = load_checkpoint(model, 'output/checkpoints/model_epoch_10.pth')
    model.load_state_dict(checkpoint['state_dict'])
    
    # 数据加载
    test_dataset = MyDataset('data/test', transform=get_transforms())
    test_loader = DataLoader(test_dataset, batch_size=64)

    # 测试模型
    test(model, test_loader)

2.7 utils.py（工具函数）

包括模型保存、加载等功能。

import torch

def save_checkpoint(model, epoch, filename):
    checkpoint = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
    }
    torch.save(checkpoint, filename)
    print(f'Model checkpoint saved to {filename}')

def load_checkpoint(model, filename):
    checkpoint = torch.load(filename)
    model.load_state_dict(checkpoint['state_dict'])
    print(f'Model checkpoint loaded from {filename}')
    return checkpoint

3. nn.Module介绍

nn.Module 是 PyTorch 中定义神经网络模型的基类。任何模型（如全连接网络、卷积神经网络、循环神经网络等）都应该继承自 nn.Module。它提供了许多关键的功能和方法，用于简化模型的定义、训练、评估以及模型参数的管理。

3.1 nn.Module 介绍

nn.Module 作为所有神经网络模型的基类，提供了以下主要功能：

• 模型层的定义：可以通过定义子类并在 __init__ 方法中定义层（如 nn.Linear, nn.Conv2d, nn.RNN 等）来创建神经网络。
• 前向传播：必须实现 forward 方法，定义模型的前向传播逻辑。
• 自动管理参数：nn.Module 会自动注册在 __init__ 方法中创建的所有层参数，确保它们可以通过 .parameters() 方法获取。
• GPU/CPU设备切换：可以使用 .to(device) 将模型和数据迁移到特定设备（如GPU）。
• 优化器与反向传播支持：nn.Module 支持与优化器结合使用，自动管理模型的梯度计算。

3.2 基本用法

3.2.1 定义一个简单的神经网络

首先，我们通过继承 nn.Module 定义一个简单的全连接神经网络（如一个两层的 MLP）：

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        # 定义层：一个输入层，两个全连接层
        self.fc1 = nn.Linear(784, 128)  # 输入784维，输出128维
        self.fc2 = nn.Linear(128, 10)   # 输入128维，输出10维（10个类别）

    def forward(self, x):
        # 前向传播
        x = F.relu(self.fc1(x))  # 使用ReLU激活函数
        x = self.fc2(x)          # 最后一层不使用激活函数
        return x

• __init__ 方法：初始化网络的层。在这里我们定义了两个全连接层（fc1 和 fc2），nn.Linear 层自动管理输入和输出大小。
• forward 方法：定义前向传播。在这里，我们使用了 ReLU 激活函数来激活第一个全连接层的输出，第二个全连接层直接输出。

3.2.2 创建模型并进行前向传播

# 创建模型实例
model = SimpleNN()

# 输入一个随机张量，假设输入是28x28的图片（展平为784维）
input_data = torch.randn(64, 784)  # 64是batch_size，784是输入特征维度

# 进行前向传播
output_data = model(input_data)

# 输出结果的形状
print(output_data.shape)  # 结果形状应该是 (64, 10)，64是batch_size，10是类别数

在上面的代码中，我们创建了 SimpleNN 类的一个实例，并输入了一个形状为 (64, 784) 的随机张量进行前向传播，输出的形状为 (64, 10)，代表每个样本有10个类别的预测值。

3.3 更复杂的网络示例

我们还可以使用 nn.Module 来构建更加复杂的神经网络模型，比如卷积神经网络（CNN）：

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # 卷积层1
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1)
        # 卷积层2
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        # 全连接层
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 卷积操作 + ReLU + 池化
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2)  # 2x2 最大池化
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)  # 2x2 最大池化
        # 扁平化（flatten）
        x = x.view(-1, 64 * 7 * 7)
        # 全连接操作
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

在这个 CNN 示例中：

• 卷积层：nn.Conv2d 定义了2个卷积层。
• 池化层：F.max_pool2d 实现了2x2的最大池化操作。
• 全连接层：nn.Linear 定义了2个全连接层，用于分类。

3.3.1 创建 CNN 模型并进行前向传播

# 创建模型实例
cnn_model = SimpleCNN()

# 输入一个28x28的灰度图像（单通道）
input_data = torch.randn(64, 1, 28, 28)  # 64是batch_size，1是通道数，28x28是图片尺寸

# 进行前向传播
output_data = cnn_model(input_data)

# 输出结果的形状
print(output_data.shape)  # 结果应该是 (64, 10)，64是batch_size，10是类别数

3.4 自动管理模型参数

nn.Module 自动管理模型中的参数，并允许我们通过 model.parameters() 来访问它们。这些参数是模型的权重和偏置。

# 获取模型的所有参数
for param in cnn_model.parameters():
    print(param.shape)

这个方法会遍历模型中的所有可训练参数（如卷积层的权重、全连接层的权重等）。

3.5 将模型移动到 GPU 或 CPU

PyTorch 使得模型在 CPU 和 GPU 之间迁移非常简单。我们只需调用 .to(device) 方法即可。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将模型移动到设备
cnn_model.to(device)

# 将输入数据也移动到相同设备
input_data = input_data.to(device)

# 进行前向传播
output_data = cnn_model(input_data)

3.6 使用 nn.Module 进行模型训练

在训练过程中，我们通常会通过定义优化器（如 Adam 或 SGD）、损失函数（如 CrossEntropyLoss）来更新模型的参数。以下是一个训练循环的简单例子：

import torch.optim as optim

# 创建模型实例
model = SimpleNN()

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 假设有训练数据（input_data 和 labels）
input_data = torch.randn(64, 784)  # 假设有64个样本，784维
labels = torch.randint(0, 10, (64,))  # 10类标签

# 训练循环
model.train()  # 设置模型为训练模式
for epoch in range(10):
    optimizer.zero_grad()  # 清除梯度
    outputs = model(input_data)  # 前向传播
    loss = criterion(outputs, labels)  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新权重
    print(f'Epoch [{epoch+1}/10], Loss: {loss.item()}')

7. 总结

nn.Module 是 PyTorch 中定义深度学习模型的核心类。它提供了：

• 模型层的自动管理和参数注册；
• 简单的前向传播实现；
• GPU/CPU设备迁移的简便方法；
• 方便的优化器与反向传播支持。

通过继承 nn.Module，你可以非常方便地定义和训练各种神经网络模型，无论是简单的全连接网络，还是复杂的卷积神经网络（CNN）、循环神经网络（RNN）等。