pytorch常用代码

在PyTorch中，有一些常用的代码片段可以帮助我们进行深度学习任务。以下是一些常见的代码片段：

1. 定义网络模型：

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)
    
    def forward(self, x):
        x = self.fc(x)
        return x

net = Net()

在这个示例中，我们定义了一个简单的神经网络模型Net，它包含一个全连接层nn.Linear。我们可以使用这个模型来进行前向传播。

2. 数据加载：

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

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __getitem__(self, index):
        x, y = self.data[index]
        return torch.tensor(x), torch.tensor(y)
    
    def __len__(self):
        return len(self.data)

dataset = CustomDataset(data)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

在这个示例中，我们定义了一个自定义的数据集CustomDataset，并使用DataLoader将数据集加载为批量大小为32的数据加载器。

3. 模型训练：

import torch
import torch.optim as optim
import torch.nn as nn

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(num_epochs):
    running_loss = 0.0
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f"Epoch {epoch+1}: Loss = {running_loss/len(dataloader)}")

在这个示例中，我们使用交叉熵损失函数nn.CrossEntropyLoss和随机梯度下降优化器optim.SGD来训练我们的模型。我们使用for循环迭代数据加载器中的批次，并在每个批次上进行前向传播、计算损失、反向传播和优化模型参数。

4. 模型保存和加载：

torch.save(net.state_dict(), 'model.pth')

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

在这个示例中，我们使用torch.save将模型的状态字典保存到文件'model.pth'中，并使用model.load_state_dict加载模型的参数。

这些是PyTorch中一些常见的代码片段，可以帮助我们进行深度学习任务。需要根据具体任务的需求进行适当的修改和调整。