在基于PyTorch的深度学习框架中,通常需要将以下内容传到GPU:
- 模型参数:神经网络模型的权重和偏置。
- 输入数据:训练和测试时的输入数据张量。
- 标签数据:对应的标签数据张量(如果有)。
- 损失函数和优化器状态:虽然损失函数和优化器本身不需要移动到GPU,但它们内部使用的数据和模型参数需要在GPU上。
步骤
- 设置设备
# 只有一块GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 指定某个GPU
device = torch.device("cuda:1")
- 将模型移动到GPU:
model = MyNeuralNet()
model.to(device)
- 将输入数据和标签移动到GPU:
inputs = inputs.to(device)
labels = labels.to(device)
- 损失函数和优化器:
损失函数不需要显式地移动到GPU,因为它会自动处理张量的位置。
优化器(optimizer)需要在模型参数移动到GPU之后定义。
显存管理
由于数据是一批一批地读进显存的,当每个批次的数据被用完之后,GPU显存中的这些数据会被释放,为下一个批次的数据腾出空间。
因此,适当降低批次大小(batch size)是一个常用的方法来节省显存。