Hello! 欢迎来到小浪云!


PyTorch在Ubuntu上如何进行模型训练


avatar
小浪云 2025-05-26 21

Ubuntu上使用pytorch进行模型训练,可以按照以下步骤进行:

安装pytorch

  1. 选择合适的安装命令: 根据你的CUDA版本(如果你有NVIDIA GPU),选择相应的PyTorch安装命令。可以在PyTorch官网找到最新的安装命令。

    # 例如,对于CUDA 11.7 pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117 
  2. 验证安装: 打开Python解释器并尝试导入PyTorch来验证安装是否成功。

    import torch print(torch.__version__) print(torch.cuda.is_available())  # 如果有GPU,应该返回True 

准备数据集

  1. 下载数据集: 可以使用PyTorch内置的数据集,或者从其他来源下载。

    from torchvision import datasets, transforms  # 定义数据转换 transform = transforms.Compose([     transforms.ToTensor(),     transforms.Normalize((0.5,), (0.5,)) ])  # 下载训练数据集 train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) 
  2. 创建数据加载器: 使用DataLoader来批量加载数据。

    from torch.utils.data import DataLoader  train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) 

定义模型

  1. 使用PyTorch的nn模块定义模型

    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, 10, kernel_size=5)         self.conv2 = nn.Conv2d(10, 20, kernel_size=5)         self.conv2_drop = nn.Dropout2d()         self.fc1 = nn.Linear(320, 50)         self.fc2 = nn.Linear(50, 10)      def forward(self, x):         x = F.relu(F.max_pool2d(self.conv1(x), 2))         x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))         x = x.view(-1, 320)         x = F.relu(self.fc1(x))         x = F.dropout(x, training=self.training)         x = self.fc2(x)         return F.log_softmax(x, dim=1) 

训练模型

  1. 初始化模型、损失函数和优化器

    model = Net() optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5) criterion = nn.CrossEntropyLoss() 
  2. 编写训练循环

    for epoch in range(10):  # 多次循环遍历数据集     running_loss = 0     for i, data in enumerate(train_loader, 0):         # 获取输入数据         inputs, labels = data          # 梯度清零         optimizer.zero_grad()          # 前向传播   反向传播   优化         outputs = model(inputs)         loss = criterion(outputs, labels)         loss.backward()         optimizer.step()          # 打印统计信息         running_loss  = loss.item()         if i % 100 == 99:    # 每100个mini-batches打印一次             print('[%d, ]] loss: %.3f' %                   (epoch   1, i   1, running_loss / 100))             running_loss = 0.0  print('Finished Training') 

保存和加载模型

  1. 保存模型

    PATH = './cifar_net.pth' torch.save(model.state_dict(), PATH) 
  2. 加载模型

    model = Net() model.load_state_dict(torch.load(PATH)) 

注意事项

  • 确保你的系统已经安装了CUDA和cuDNN,如果你打算使用GPU进行训练。
  • 根据你的具体需求调整模型结构、损失函数和优化器。
  • 监控训练过程中的损失和准确率,以便及时调整超参数。

以上步骤提供了一个基本的框架,你可以根据自己的项目需求进行调整和扩展

相关阅读