当前位置 博文首页 > pytorch 如何在GPU上训练

    pytorch 如何在GPU上训练

    作者:陌生的天花板 时间:2021-08-12 17:43

    1.网络模型转移到CUDA上

    net = AlexNet()
    net.cuda()#转移到CUDA上

    2.将loss转移到CUDA上

    criterion = nn.CrossEntropyLoss()
    criterion = criterion.cuda()

    这一步不做也可以,因为loss是根据out、label算出来的

    loss = criterion(out, label)

    只要out、label在CUDA上,loss自然也在CUDA上了,但是发现不转移到CUDA上准确率竟然降低了1%

    3.将数据集转移到CUDA上

    这里要解释一下数据集使用方法

    #download the dataset
    train_set = CIFAR10("./data_cifar10", train=True, transform=data_tf, download=True)
    train_data = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

    dataset是把所有的input,label都制作成了一个大的多维数组

    dataloader是在这个大的多维数组里采样制作成batch,用这些batch来训练

        for im, label in train_data:
            i = i + 1
            im = im.cuda()#把数据迁移到CUDA上
            im = Variable(im)#把数据放到Variable里
            label = label.cuda()
            label =Variable(label)
            out = net(im)#the output should have the size of (N,10)

    遍历batch的时候,首先要把拿出来的Image、label都转移到CUDA上,这样接下来的计算都是在CUDA上了

    开始的时候只在转成Variable以后才迁移到CUDA上,这样在网络传播过程中就数据不是在CUDA上了,所以一直报错

    训练网络时指定gpu显卡

    查看有哪些可用的gpu

    nvidia -smi

    实时查看gpu信息1代表每1秒刷新一次

    watch -n -1 nvidia -smi

    指定使用的gpu

    import os
    # 使用第一张与第三张GPU卡
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,3"

    以上为个人经验,希望能给大家一个参考,也希望大家多多支持站长博客。

    jsjbwy