在训练模型的时候,不可避免的要使用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)