在Ubuntu上使用PyTorch训练模型的步骤如下:
首先,你需要安装PyTorch。你可以根据你的CUDA版本选择合适的安装命令。以下是一些常用的安装命令:
pip install torch torchvision torchaudio
如果你需要GPU支持,可以安装带有CUDA支持的版本:
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
请根据你的CUDA版本选择合适的URL(例如,cu113对应CUDA 11.3)。
如果你使用Anaconda,可以使用以下命令安装:
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
同样,请根据你的CUDA版本选择合适的cudatoolkit。
准备你的训练数据集。你可以使用PyTorch提供的torchvision.datasets
模块来加载常见的数据集,例如MNIST、CIFAR-10等。
import torchvision.transforms as transforms
import torchvision.datasets as datasets
# 定义数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
# 加载训练数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
定义你的神经网络模型。你可以使用PyTorch提供的预定义模型,也可以自己定义。
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 3)
self.fc1 = nn.Linear(6 * 6 * 6, 10)
def forward(self, x):
x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
x = x.view(-1, self.num_flat_features(x))
x = self.fc1(x)
return x
def num_flat_features(self, x):
size = x.size()[1:] # all dimensions except the batch dimension
num_features = 1
for s in size:
num_features *= s
return num_features
net = Net()
定义损失函数和优化器。
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
编写训练循环来训练模型。
for epoch in range(5): # 多次循环遍历数据集
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
# 获取输入数据
inputs, labels = data
# 梯度清零
optimizer.zero_grad()
# 前向传播 + 反向传播 + 优化
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 打印统计信息
running_loss += loss.item()
if i % 100 == 99: # 每100个mini-batches打印一次
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 100))
running_loss = 0.0
print('Finished Training')
训练完成后,你可以保存模型以便以后使用。
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)
你可以加载保存的模型并进行预测。
net = Net()
net.load_state_dict(torch.load(PATH))
dataiter = iter(train_loader)
images, labels = dataiter.next()
# 打印图像和标签
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))
以上步骤涵盖了在Ubuntu上使用PyTorch训练模型的基本流程。根据具体需求,你可能需要调整数据预处理、模型结构、损失函数和优化器等部分。