Pytorch实验常用代码段汇总-猿码集

1. Pytorch实验常用代码段汇总

在深度学习实验中，常用Pytorch实现模型训练和评估。为了方便实验，在这里整理了一些Pytorch实验中常用的代码段。本文将从模型定义、数据处理、训练和测试等几个方面来介绍。

2. 模型定义

2.1 定义模型

在Pytorch中，我们可以通过继承torch.nn.Module类来自定义模型，在__init__函数中定义网络结构，在forward函数中实现前向计算。下面是一个简单的卷积神经网络的定义：

import torch.nn as nn
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.fc = nn.Linear(32 * 7 * 7, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 32 * 7 * 7)
        x = self.fc(x)
        return x

在上面的代码中，我们定义了一个包含两个卷积层和一个全连接层的卷积神经网络。在初始化函数__init__中，我们定义了各层的参数（如卷积核大小、步长、padding等），在forward函数中实现了前向传播。其中，F.relu表示激活函数ReLU，self.pool表示池化层，x.view表示重塑张量维度。

2.2 模型初始化

在使用Pytorch定义模型后，需要初始化模型参数。常用的初始化方式有Xavier初始化和Kaiming初始化，其中Xavier适用于sigmoid等饱和的激活函数，Kaiming适用于ReLU等不饱和的激活函数。下面以Xavier初始化为例：

import torch.nn.init as init
def init_weights(net):
    for m in net.modules():
        if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
            init.xavier_normal(m.weight.data)
            if m.bias is not None:
                m.bias.data.zero_()

在上面的代码中，我们定义了一个init_weights函数，对所有卷积层和全连接层进行Xavier初始化。其中，init.xavier_normal函数表示使用Xavier初始化方法初始化权重，zero_函数表示初始化偏移为0。

2.3 模型加载与保存

在训练模型过程中，我们需要保存模型以便后续使用。Pytorch提供了保存和加载模型的函数，如下所示：

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

在上面的代码中，我们使用torch.save函数保存模型参数，使用model.load_state_dict函数加载模型参数。

3. 数据处理

3.1 数据加载

在Pytorch中，可以通过Dataset和DataLoader实现数据加载。Dataset用于准备数据，DataLoader用于加载数据并生成batch。下面是一个MNIST数据集的加载示例：

from torchvision import datasets, transforms
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=True)

在上面的代码中，我们使用transforms.Compose函数定义数据变换（如将图片转换为Tensor并进行标准化）。然后使用datasets.MNIST函数加载数据集，并使用DataLoader生成batch，shuffle表示是否随机打乱数据。

3.2 自定义数据集

除了使用内置的数据集，我们也可以自定义数据集。自定义数据集时，需要继承torch.utils.data.Dataset类，并实现__len__和__getitem__函数。下面是一个自定义的二分类数据集：

import torch.utils.data as data
class BinaryDataset(data.Dataset):
    def __init__(self, data, targets):
        self.data = data
        self.targets = targets
    def __getitem__(self, index):
        x, y = self.data[index], self.targets[index]
        return torch.Tensor(x), torch.LongTensor([y])
    def __len__(self):
        return len(self.targets)

在上面的代码中，我们定义了一个BinaryDataset类，包含数据和标签。__getitem__函数用于获取指定索引的数据，并将数据转换为tensor类型。__len__函数返回数据集大小。

4. 训练和测试

4.1 梯度更新

在训练模型时，需要反向传播计算梯度并更新模型参数。Pytorch使用自动微分机制自动计算梯度，并使用optimezer优化器更新模型参数。下面是一个梯度更新的示例：

import torch.optim as optim
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(epochs):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

在上面的代码中，我们使用torch.optim.SGD函数定义优化器为随机梯度下降方法，learning_rate为0.01，momentum为0.9，在每个epoch中，对每个batch计算损失和反向传播更新梯度。

4.2 模型评估

在测试集上，我们需要评估模型的表现。通常使用准确率、召回率、F1-score等指标进行评估。下面是一个测试模型的示例：

def test_model(model, test_loader):
    correct = 0
    total = 0
    model.eval()
    with torch.no_grad():
        for inputs, targets in test_loader:
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += targets.size(0)
            correct += (predicted == targets).sum().item()
    print('Accuracy on test set: %d %%' % (100 * correct / total))
test_model(model, test_loader)

在上面的代码中，我们使用torch.no_grad函数表示不计算梯度。在每个batch中，计算模型输出并与目标值进行比较，统计预测正确的样本数，最后计算准确率。

4.3 模型预测

在使用训练好的模型对新数据进行预测时，需要将输入数据转换为tensor类型，并对模型输出进行后处理。下面是一个模型预测的示例：

import torch.nn.functional as F
def predict(model, input):
    model.eval()
    input = torch.Tensor(input)
    output = model(input)
    output = F.softmax(output / temperature, dim=1)
    pred = output.argmax(dim=1).item()
    return pred

在上面的代码中，我们定义了一个predict函数，用于对输入数据进行预测。首先将输入数据转换为tensor类型，然后计算模型输出，进行softmax处理并选取概率最大的类别。

5. 总结

本文介绍了一些Pytorch实验中常用的代码段，包括模型定义、数据处理、训练和测试等方面。通过了解这些代码段，可以更加高效地进行Pytorch实验。

Pytorch实验常用代码段汇总