1.使用cuda方法在模型，数据加载与损失函数下进行gpu加速

class yu(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super(yu, self).__init__(*args, **kwargs)

        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10),
        )
    def forward(self, x):
        x = self.model(x)
        return x

hh=yu()
if torch.cuda.is_available():   
    hh=hh.cuda()#对网络模型用cuda方法
loss = nn.CrossEntropyLoss()
loss=loss.cuda()#对损失函数用cuda方法

for data in train_dataloader:
    imgs,target =data
    imgs = imgs.cuda()
    target =target.cuda()#对数据集进行cuda方法

1.1没有GPU的同学可以使用google colab 

2.也可以使用torch.device()进行设备的选取

#定义训练的设备
device =torch.device("cpu")
class yu(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super(yu, self).__init__(*args, **kwargs)

        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10),
        )
    def forward(self, x):
        x = self.model(x)
        return x

hh=yu()
hh.to(device)#将模型加载到我们定义选择的芯片设备

2.1 如果有多块卡，可以写成

device = torch.device("cuda:0")#这里0指我们用的是第一块gpu