阅读量:0
在Ubuntu下编写PyTorch训练脚本涉及几个步骤，包括环境设置、数据准备、模型定义、损失函数和优化器的选择，以及训练循环的实现。以下是一个简单的示例，展示了如何编写一个基本的PyTorch训练脚本。
步骤1：安装PyTorch
首先，确保你已经安装了PyTorch。你可以参考PyTorch官网的安装指南来安装适合你系统的版本。
步骤2：编写训练脚本
创建一个新的Python文件，例如train.py，并编写以下代码：
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 定义一个简单的神经网络
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.fc1 = nn.Linear(784, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = x.view(-1, 784) # Flatten the input
x = self.relu(self.fc1(x))
x = self.relu(self.fc2(x))
x = self.fc3(x)
return x
# 数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
# 初始化模型、损失函数和优化器
model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# 训练模型
num_epochs = 5
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')
print('Finished Training')
代码解释
-
导入必要的库：
torch和torch.nn用于构建和训练神经网络。torch.optim用于定义优化器。torchvision用于加载和预处理数据。DataLoader用于批量加载数据。
-
定义模型：
SimpleNet是一个简单的全连接神经网络，包含三个全连接层和一个ReLU激活函数。
-
数据预处理：
- 使用
transforms.Compose组合多个数据预处理操作。 ToTensor将图像转换为PyTorch张量。Normalize对图像进行标准化。
- 使用
-
加载数据集：
- 使用
datasets.MNIST加载MNIST数据集，并应用预处理操作。 - 使用
DataLoader批量加载数据，并设置shuffle=True以打乱数据顺序。
- 使用
-
初始化模型、损失函数和优化器：
SimpleNet实例化模型。CrossEntropyLoss作为损失函数。SGD作为优化器，设置学习率和动量。
-
训练模型：
- 使用嵌套的
for循环进行训练。 - 外层循环遍历每个epoch。
- 内层循环遍历每个batch的数据，计算损失并进行反向传播和参数更新。
- 打印每个epoch的损失。
- 使用嵌套的
运行脚本
在终端中运行以下命令来执行训练脚本：
python train.py
这个示例展示了如何在Ubuntu下编写一个基本的PyTorch训练脚本。你可以根据具体需求扩展和修改这个脚本，例如添加验证步骤、保存模型、使用GPU加速等。
以上就是关于“Ubuntu下PyTorch训练脚本怎么写”的相关介绍，筋斗云是国内较早的云主机应用的服务商，拥有10余年行业经验，提供丰富的云服务器、租用服务器等相关产品服务。云服务器资源弹性伸缩，主机vCPU、内存性能强悍、超高I/O速度、故障秒级恢复；电子化备案，提交快速，专业团队7×24小时服务支持！
简单好用、高性价比云服务器租用链接：https://www.jindouyun.cn/product/cvm