使用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
一定要额外赋值,为了方便记忆还是都选择使用赋值的方式,但要能看懂别人不进行赋值的代码