PyTorch中的CPU和GPU代码实现详解

PyTorch中的CPU和GPU代码实现详解

在深度学习的开发过程中，计算资源的高效利用是至关重要的。PyTorch作为一种流行的深度学习框架，支持使用CPU和GPU进行模型训练和推理。相较于CPU，GPU由于其强大的并行计算能力，能够显著加速深度学习任务。然而，将PyTorch代码从CPU版本迁移到GPU版本需要进行一些额外的代码修改。本文将详细介绍如何在PyTorch中编写支持CPU和GPU的代码，以及需要特别注意的事项。

1. 安装PyTorch

首先，确保你已经安装了支持GPU的PyTorch版本。如果还没有安装，可以参考以下命令进行安装：

# For CUDA 11.1
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu111

2. 编写支持CPU和GPU的PyTorch代码

2.1模型定义

定义模型的代码在CPU和GPU版本中基本一致。但是，我们需要确保模型可以在GPU上运行。

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        return self.fc(x)

model = SimpleNN()

2.2数据加载

数据加载部分对于CPU和GPU是相同的。使用DataLoader类加载数据：

在这里插入代码片from torchvision import datasets, transforms

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,),