阅读量:4
在Ubuntu上训练PyTorch模型的步骤如下:
1. 安装Python和PyTorch
首先,确保你的Ubuntu系统已经安装了Python。推荐使用Python 3.6或更高版本。
安装Python
sudo apt update
sudo apt install python3 python3-pip
安装PyTorch
访问PyTorch官网,根据你的CUDA版本选择合适的安装命令。以下是使用pip安装的示例:
-
CPU版本:
pip3 install torch torchvision torchaudio -
GPU版本(CUDA 11.7):
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
2. 准备数据集
你需要准备用于训练的数据集。常见的做法是将数据集放在一个目录中,并使用torchvision.transforms进行预处理。
示例:使用CIFAR-10数据集
import torchvision.transforms as transforms
import torchvision.datasets as datasets
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, padding=4),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)
testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)
3. 定义模型
你可以使用PyTorch提供的预定义模型,或者自己定义一个模型。
示例:使用ResNet18
import torch.nn as nn
import torchvision.models as models
model = models.resnet18(pretrained=False)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10) # CIFAR-10有10个类别
4. 定义损失函数和优化器
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
5. 训练模型
for epoch in range(10): # 多个epoch
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 200 == 199: # 每200个batch打印一次
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 200:.3f}')
running_loss = 0.0
print('Finished Training')
6. 测试模型
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')
7. 保存模型
PATH = './cifar_net.pth'
torch.save(model.state_dict(), PATH)
8. 加载模型
model = models.resnet18(pretrained=False)
model.fc = nn.Linear(num_ftrs, 10)
model.load_state_dict(torch.load(PATH))
通过以上步骤,你可以在Ubuntu上训练一个PyTorch模型。根据具体需求,你可以调整数据预处理、模型结构、损失函数和优化器等参数。
以上就是关于“Ubuntu上PyTorch模型怎么训练”的相关介绍,筋斗云是国内较早的云主机应用的服务商,拥有10余年行业经验,提供丰富的云服务器、租用服务器等相关产品服务。云服务器资源弹性伸缩,主机vCPU、内存性能强悍、超高I/O速度、故障秒级恢复;电子化备案,提交快速,专业团队7×24小时服务支持!
简单好用、高性价比云服务器租用链接:https://www.jindouyun.cn/product/cvm