Pytorch深度学习 - Class Model()和训练函数

下面是一个示例代码，展示了如何使用PyTorch来定义一个深度学习模型并编写训练函数。

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
    
    def forward(self, x):
        x = self.fc1(x)
        x = nn.ReLU()(x)
        x = self.fc2(x)
        return x

def train(model, dataloader, criterion, optimizer):
    model.train()
    
    for inputs, labels in dataloader:
        # 清零梯度
        optimizer.zero_grad()
        
        # 前向传播
        outputs = model(inputs)
        
        # 计算损失
        loss = criterion(outputs, labels)
        
        # 反向传播
        loss.backward()
        
        # 更新权重
        optimizer.step()

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

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 创建数据加载器（这里假设已经准备好了训练数据）
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 进行训练
train(model, train_dataloader, criterion, optimizer)

在这个示例中，我们首先定义了一个自定义的深度学习模型MyModel，该模型包含两个线性层和一个ReLU激活函数。然后，我们定义了一个训练函数train，该函数接受模型、数据加载器、损失函数和优化器作为输入，并在训练过程中执行前向传播、损失计算、反向传播和权重更新的步骤。在主代码中，我们创建了模型实例、定义了损失函数和优化器，并使用数据加载器进行训练。请注意，这里的数据加载器train_dataloader是一个假设已经准备好的训练数据的示例，你需要根据自己的数据进行相应的调整。