PyTorch 模型训练代码示例 | 深度学习教程

下面是一个使用 PyTorch 训练模型的示例代码：\n\npython\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\n\n# 定义模型\nclass MyModel(nn.Module):\n def __init__(self):\n super(MyModel, self).__init__()\n self.fc = nn.Linear(10, 1) # 输入维度为10，输出维度为1\n\n def forward(self, x):\n x = self.fc(x)\n return x\n\n# 准备训练数据\ntrain_data = torch.randn((100, 10)) # 100个样本，每个样本有10个特征\ntrain_labels = torch.randn((100, 1)) # 100个样本，每个样本有1个标签\n\n# 初始化模型\nmodel = MyModel()\n\n# 定义损失函数和优化器\ncriterion = nn.MSELoss()\noptimizer = optim.SGD(model.parameters(), lr=0.01)\n\n# 进行训练\nnum_epochs = 10\nfor epoch in range(num_epochs):\n model.train() # 设置模型为训练模式\n optimizer.zero_grad() # 梯度清零\n\n # 前向传播\n outputs = model(train_data)\n loss = criterion(outputs, train_labels)\n\n # 反向传播和优化\n loss.backward()\n optimizer.step()\n\n # 打印训练信息\n print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')\n\n# 保存模型\ntorch.save(model.state_dict(), 'model.pth')\n\n\n这段代码首先定义了一个简单的模型MyModel，包含一个全连接层。然后准备训练数据，使用随机生成的数据作为训练集。接着定义了损失函数和优化器，这里使用均方误差作为损失函数，使用随机梯度下降作为优化器。然后使用循环进行训练，每个epoch中进行前向传播、计算损失、反向传播和优化。最后保存训练好的模型参数。