Pytorch多卡训练

首先查询是否支持 CUDA 训练 进行指定设备

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

使用 GPU 训练时 需要增加以下代码

import torch

import torch.nn as nn

*# 如果你有多个 GPU，可以使用 torch.nn.DataParallel 来并行训练 *

if torch.cuda.device_count() > 1:

print(f"Let's use {torch.cuda.device_count()} GPUs!")

model = nn.DataParallel(model)

设置完成后 启动训练时会显示

显卡占用低的 或者 训练时间较慢时的调试

1. 显存占用低时 需要增加 batch_size

2. 显卡核心占用低时需要按照实际情况调节 DataLoader 方法内的 nums_worker