Pytorch使用GPU训练

1. 查看是否有可用的GPU

use_gpu = torch.cuda.is_available() 

use_gpu 为True则表示有GPU可用

2. 将网络的模型和loss函数放到GPU中

loss_fc = nn.MSELoss()
    # While constructing the network, transfer the model to GPU (pytorch)
    if (use_gpu):
        model = torch.nn.DataParallel(model, device_ids=[0]).cuda()
        # model = model.cuda()
        loss_fc = loss_fc.cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr=learn_rate, betas=(0.9, 0.99), weight_decay=w_decay)

3. 将训练的数据数据放入GPU中，数据类型需为Tensor

        for i, data in enumerate(data_loader):
            train_x, train_y = data
            # While traning, transfer the data to GPU
            if (use_gpu):
                train_x, train_y = train_x.cuda(), train_y.cuda()

4. 最后从GPU中取出需要最后输出的数据

    if (use_gpu):
        loss,prediction_y= loss.cpu(),prediction_y.cpu()