一般来说,我们可以对:
网络的模型,损失函数,数据集,利用GPU进行运算也就是调用cuda.
第一种方式:
my_model = Model()
if torch.cuda.is_available():
my_model = my_model.cuda()
loss_cross = nn.CrossEntropyLoss()
if torch.cuda.is_available():
loss_cross = loss_cross.cuda()
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
直接调用cuda.
第二种形式:
先创建一个你指定的设备divice
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
然后将要在gpu上训练的代码块添加到该设备上
my_model = Model()
my_model.to(device)
loss_cross = nn.CrossEntropyLoss()
loss_cross.to(device)
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
第二种方式更常用,因为它可以指定你用哪块显卡,只需要在"cuda"后面进行指定,如"cuda:0","cuda1",看上去更专业吧