pytorch学习（三）cpu-gpu训练

在训练模型的时候，不可避免的要使用GPU进行加速，但是我们一般加载或者创建生成的数据都是处于CPU上，怎么把数据加载到GPU上呢？

1.首先看变量的CPU和GPU转换

import torch


tensor = torch.randn(3,3)
bTensor = type(tensor) == torch.Tensor
print(bTensor)
print("tensor is on ", tensor.device)
#数据转到GPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
if torch.cuda.is_available():
    tensor = tensor.to(device)
    print("tensor is on ",tensor.device)
#数据转到CPU
# if tensor.device == 'cuda:0':
    tensor = tensor.to(torch.device("cpu"))
    print("tensor is on", tensor.device)
# if tensor.device == "cpu":
    tensor = tensor.to(torch.device("cuda:0"))
    print("tensor is on", tensor.device)

2.训练过程中如何转换呢？

主要是数据和模型的转换

#1.第一个地方
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

#2.把模型放到GPU上
class MinistNet(nn.Module):
    def __init__(self):
        super().__init__()
        # self.flat = nn.Flatten()
        self.conv1 = nn.Conv2d(1,1,3,1,1)
        self.hideLayer1 = nn.Linear(28*28,256)
        self.hideLayer2 = nn.Linear(256,10)
    def forward(self,x):
        x= self.conv1(x)
        x = x.view(-1,28*28)
        x = self.hideLayer1(x)
        x = torch.sigmoid(x)
        x = self.hideLayer2(x)
        # x = nn.Sigmoid(x)
        return x

model = MinistNet()
model = model.to(device)  #把模型加载到GPU上，
# 或者使用 model.to(device)直接就把模型加载到GPU上

#3.把数据放到GPU上
        for x,y in train_loader:
            # print(epoch)
            # print(x.shape)
            # print(y.shape)
            x = x.to('cuda') #直接使用x.to(device)是不能把x数据加载到GPU上，必须使用x=x.to(device)
            y = y.to('cuda')
            A = x.device
            B = y.device
            pred_y = model(x)

初次学习需要注意的是 ：直接使用x.to(device)是不能把x数据加载到GPU上，必须使用x=x.to(device)