UV-Net的复现，能不能用GPU

Question

你好！我今天正在看你复现UV-Net那篇文章，在我复现的时候无法调用GPU训练的很慢。你在环境配置中写道“上面那种安装有很多没有指定其版本故若要进行GPU版本的pytorch运算可能不太行。若你只是想能跑或者自己本身的CUDA版本够高的话运行起来是没问题的。”请问如何才能调用GPU跑呢？

Answer 1

以下回复参考：皆我百晓生、券券喵儿等免费微信小程序作答:

您好！关于如何调用GPU进行训练，您需要确保以下几点：

1. 确保您的计算机上已经安装了NVIDIA的CUDA工具包和cuDNN库。您可以从NVIDIA官网下载并安装它们。

2. 安装支持GPU的PyTorch版本。您可以使用以下命令安装：

pip install torch torchvision -f https://download.pytorch.org/whl/cu102/torch_stable.html

请根据您的CUDA版本替换cu102。例如，如果您的CUDA版本是10.1，则应使用cu101。

1. 在您的代码中，确保使用torch.cuda.is_available()检查GPU是否可用，并使用torch.cuda.device(0)将GPU设备设置为默认设备。例如：

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

1. 在创建模型、优化器和数据加载器时，确保将它们放在GPU设备上。例如：

model = UV_Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

1. 在训练循环中，确保将输入数据和模型放在GPU设备上。例如：

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(dataloader):
        images = images.to(device)
        labels = labels.to(device)
        # ...

按照以上步骤操作后，您应该可以使用GPU进行训练了。祝您复现UV-Net顺利！