利用GPU进行训练

一般来说，我们可以对：

网络的模型，损失函数，数据集，利用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"，看上去更专业吧

版权声明：本文为CSDN博主「zmnhlg」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/zmnhlg/article/details/121629003