Ubuntu上PyTorch训练模型步骤 - 问答

在Ubuntu上使用PyTorch训练模型的步骤如下：

1. 安装Python和pip

确保你的系统上已经安装了Python和pip。Ubuntu通常自带Python，但可能需要更新到最新版本。

sudo apt update
sudo apt install python3 python3-pip

2. 创建虚拟环境（可选）

为了隔离项目依赖，建议创建一个虚拟环境。

sudo apt install python3-venv
python3 -m venv myenv
source myenv/bin/activate

3. 安装PyTorch

根据你的CUDA版本选择合适的PyTorch安装命令。可以在PyTorch官网找到最新的安装命令。

例如，如果你使用CUDA 11.7：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117

如果不使用CUDA：

pip install torch torchvision torchaudio

4. 安装其他依赖库

根据你的项目需求，安装其他必要的库。例如，如果你使用OpenCV：

pip install opencv-python

5. 准备数据集

将你的数据集准备好，并放在一个合适的位置。你可以使用torchvision.datasets中的数据集，或者自己编写数据加载器。

6. 编写训练脚本

创建一个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 SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc = nn.Linear(32 * 14 * 14, 10)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = x.view(-1, 32 * 14 * 14)
        x = self.fc(x)
        return x

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 初始化模型、损失函数和优化器
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
for epoch in range(5):
    running_loss = 0.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()
    print(f'Epoch {epoch + 1}, Loss: {running_loss / (i + 1)}')

print('Finished Training')

7. 运行训练脚本

在终端中运行你的训练脚本：

python train.py

8. 监控训练过程

你可以使用TensorBoard来监控训练过程。首先安装TensorBoard：

pip install tensorboard

然后在训练脚本中添加TensorBoard的代码：

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/simple_cnn_experiment')

for epoch in range(5):
    running_loss = 0.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()
    writer.add_scalar('training loss', running_loss / (i + 1), epoch)
    print(f'Epoch {epoch + 1}, Loss: {running_loss / (i + 1)}')

writer.close()

最后，在终端中启动TensorBoard：

tensorboard --logdir=runs

然后在浏览器中打开http://localhost:6006查看训练过程。

9. 保存和加载模型

训练完成后，你可以保存模型以便以后使用：

torch.save(model.state_dict(), 'model.pth')

加载模型：

model = SimpleCNN()
model.load_state_dict(torch.load('model.pth'))
model.eval()

以上就是在Ubuntu上使用PyTorch训练模型的基本步骤。根据你的具体需求，可能需要进行一些调整和优化。

0 赞

0 踩