Pytorch调用GPU训练两种方法

方法一 .cuda()

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

1.  
   # 将网络模型在gpu上训练
2.  
   model = Model()
3.  
   if torch.cuda.is_available():
4.  
   model = model.cuda()
5.  
    
6.  
   # 损失函数在gpu上训练
7.  
   loss_fn = nn.CrossEntropyLoss()
8.  
   if torch.cuda.is_available():
9.  
   loss_fn = loss_fn.cuda()
10.  
     
11.  
    # 数据在gpu上训练
12.  
    for data in dataloader:
13.  
    imgs, targets = data
14.  
    if torch.cuda.is_available():
15.  
    imgs = imgs.cuda()
16.  
    targets = targets.cuda()

方法二 .to(device)

指定 训练的设备

1.  
   device = torch.device("cpu") # 使用cpu训练
2.  
   device = torch.device("cuda") # 使用gpu训练
3.  
   device = torch.device("cuda:0") # 当电脑中有多张显卡时，使用第一张显卡
4.  
   device = torch.device("cuda:1") # 当电脑中有多张显卡时，使用第二张显卡
5.  
    
6.  
   #单卡推荐
7.  
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

使用 GPU 训练

1.  
   model = model.to(device)
2.  
    
3.  
   loss_fn = loss_fn.to(device)
4.  
    
5.  
   for data in train_dataloader:
6.  
   imgs, targets = data
7.  
   imgs = imgs.to(device)
8.  
   targets = targets.to(device)

查看GPU使用情况

nvidia-smi.exe -l 5

这篇好文章是转载于：学新通技术网