30. 使用GPU进行模型的训练(二) to()方法

使用GPU进行模型的训练(二) to()方法

1. 使用GPU要注意哪些问题

• 如果我们使用GPU训练模型，需要同时对 模型、数据、损失函数 同时指定使用GPU模型，否则就会存在运行问题

2. to()方法使用GPU

• 在使用to()方法之前，首先要定义使用的设别，为了保证代码的健壮性可以在多个系统中运行，可以在指定设备时增加如下的条件选择：

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

• to()的使用与 cuda() 的使用格式是一样的，下面给出对于模型、损失函数、数据进行 to()方法的使用代码

  # 模型
  model = Model()
  model = model.to(device)
  
  # 损失函数
  loss_fn = nn.CrossEntropyLoss()
  loss_fn = loss_fn.to(device)
  
  # 数据
  imgs = imgs.to(device)
  targets = targets.to(device)
  

  需要明确的是，如果要使用GPU，就一定要对模型、损失函数、数据(data && targets) 同时都使用 .to() 方法

Tips： Tips: 对于模型和损失函数，直接使用 model.to() loss_fn.to() 方法也是可以的，不需要进行额外的赋值操作，但是数据 imgs, targets 一定要额外赋值，为了方便记忆还是都选择使用赋值的方式，但要能看懂别人不进行赋值的代码