pytorch 训练模型代码

下面是一个使用PyTorch训练模型的示例代码：

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

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 1)  # 输入维度为10，输出维度为1

    def forward(self, x):
        x = self.fc(x)
        return x

# 准备训练数据
train_data = torch.randn((100, 10))  # 100个样本，每个样本有10个特征
train_labels = torch.randn((100, 1))  # 100个样本，每个样本有1个标签

# 初始化模型
model = MyModel()

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

# 进行训练
num_epochs = 10
for epoch in range(num_epochs):
    model.train()  # 设置模型为训练模式
    optimizer.zero_grad()  # 梯度清零

    # 前向传播
    outputs = model(train_data)
    loss = criterion(outputs, train_labels)

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

    # 打印训练信息
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

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

这段代码首先定义了一个简单的模型MyModel，包含一个全连接层。然后准备训练数据，使用随机生成的数据作为训练集。接着定义了损失函数和优化器，这里使用均方误差作为损失函数，使用随机梯度下降作为优化器。然后使用循环进行训练，每个epoch中进行前向传播、计算损失、反向传播和优化。最后保存训练好的模型参数