PyTorch GPU训练：两种方法详解

方法一 .cuda()

我们可以通过对网络模型，数据，损失函数这三种变量调用 .cuda() 来在GPU上进行训练。

# 将网络模型在gpu上训练
model = Model()
if torch.cuda.is_available():
	model = model.cuda()

# 损失函数在gpu上训练
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():	
	loss_fn = loss_fn.cuda()

# 数据在gpu上训练
for data in dataloader:                        
	imgs, targets = data
    if torch.cuda.is_available():
        imgs = imgs.cuda()
        targets = targets.cuda()

方法二 .to(device)

指定 训练的设备

device = torch.device("cpu")	# 使用cpu训练
device = torch.device("cuda")	# 使用gpu训练 
device = torch.device("cuda:0")	# 当电脑中有多张显卡时，使用第一张显卡
device = torch.device("cuda:1")	# 当电脑中有多张显卡时，使用第二张显卡

#单卡推荐
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

使用 GPU 训练

model = model.to(device)

loss_fn = loss_fn.to(device)

for data in train_dataloader:
    imgs, targets = data
    imgs = imgs.to(device)
    targets = targets.to(device)

查看GPU使用情况

nvidia-smi.exe -l 5