Python Pytorch包详解

PyTorch 包详解

PyTorch 是一个基于 Torch 的开源机器学习框架，由 Facebook 的人工智能研究团队开发。它提供了两个核心功能：

强大的 GPU 加速张量计算基于自动微分系统的深度神经网络

一、核心组件

1. Tensor（张量）

PyTorch 的基础数据结构，类似于 NumPy 的 ndarray，但支持 GPU 加速。

import torch

# 创建张量
x = torch.tensor([1, 2, 3])
x = torch.randn(2, 3)  # 随机张量
x = torch.zeros(3, 4)  # 零张量
x = torch.ones(2, 2)   # 单位张量

# 张量操作
y = x + 2
z = torch.mm(x, y.T)  # 矩阵乘法

2. Autograd（自动微分）

PyTorch 的核心特性，自动计算梯度。

# 创建需要梯度的张量
x = torch.tensor([1., 2., 3.], requires_grad=True)

# 进行计算
y = x ** 2
z = y.sum()

# 反向传播计算梯度
z.backward()
print(x.grad)  # 输出: tensor([2., 4., 6.])

二、神经网络模块

1. torch.nn 模块

构建神经网络的核心模块。

import torch.nn as nn

# 定义网络结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 128)  # 全连接层
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()           # 激活函数
        self.dropout = nn.Dropout(0.2)  # Dropout层

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 常用层类型
conv = nn.Conv2d(1, 32, 3)    # 卷积层
pool = nn.MaxPool2d(2)        # 池化层
lstm = nn.LSTM(10, 20, 2)     # LSTM层
batch_norm = nn.BatchNorm2d(3) # 批量归一化

2. 激活函数

relu = nn.ReLU()
sigmoid = nn.Sigmoid()
tanh = nn.Tanh()
leaky_relu = nn.LeakyReLU(0.01)
softmax = nn.Softmax(dim=1)

三、优化器

import torch.optim as optim

model = Net()
criterion = nn.CrossEntropyLoss()  # 损失函数

# 各种优化器
optimizer = optim.SGD(model.parameters(), lr=0.01)
optimizer = optim.Adam(model.parameters(), lr=0.001)
optimizer = optim.RMSprop(model.parameters(), lr=0.01)
optimizer = optim.Adagrad(model.parameters(), lr=0.01)

# 学习率调度器
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

四、数据加载与处理

1. torch.utils.data

from torch.utils.data import DataLoader, Dataset, TensorDataset

# 自定义数据集
class CustomDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 创建数据加载器
dataset = CustomDataset(x, y)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

2. torchvision（计算机视觉专用）

import torchvision
import torchvision.transforms as transforms

# 数据转换
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225])
])

# 加载内置数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                       download=True, transform=transform)

五、模型训练示例

def train(model, dataloader, criterion, optimizer, epochs=10):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for batch_idx, (data, target) in enumerate(dataloader):
            optimizer.zero_grad()  # 梯度清零

            # 前向传播
            output = model(data)
            loss = criterion(output, target)

            # 反向传播
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f'Epoch {epoch+1}, Loss: {running_loss/len(dataloader):.4f}')

六、GPU 加速

# 检查GPU可用性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将模型和数据移动到GPU
model = Net().to(device)
data = data.to(device)
target = target.to(device)

# 多GPU训练
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

七、保存和加载模型

# 保存模型
torch.save(model.state_dict(), 'model.pth')

# 加载模型
model.load_state_dict(torch.load('model.pth'))

# 保存整个模型
torch.save(model, 'model_full.pth')

# 保存训练检查点
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, 'checkpoint.pth')

八、高级特性

1. 混合精度训练

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

with autocast():
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

2. 分布式训练

import torch.distributed as dist
import torch.multiprocessing as mp

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

3. TorchScript（模型部署）

# 转换为 TorchScript
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')

# 跟踪模型
traced_model = torch.jit.trace(model, example_input)

九、实用工具

1. 可视化（需要安装 tensorboard）

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter()
writer.add_graph(model, input_tensor)
writer.add_scalar('Loss/train', loss, epoch)

2. 性能分析

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ]
) as prof:
    model(input)
print(prof.key_averages().table())

十、PyTorch Lightning（简化训练的高级框架）

import pytorch_lightning as pl

class LitModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = Net()
        self.criterion = nn.CrossEntropyLoss()

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = self.criterion(y_hat, y)
        return loss

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=0.001)

# 训练
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, dataloader)