背景
在电脑上运行深度学习模型,需要设置设备在cup还是GPU上运行,GPU运行速度是显著高于CPU的,但对于在不同设备上保存的变量在运行时很容易发生错误。因此需要注意
例子
from sklearn.metricsimport f1_scoreimport torchimport torch.nnas nnfrom torch.nnimport Moduleimport torch.nn.functionalas Fimport mathimport torch.optimas optimimport numpyas npclassGraphConvolution(Module):"""
A Graph Convolution Layer (GCN)
"""def__init__(self, in_features, out_features, bias=True):super(GraphConvolution, self).__init__()
self.in_features= in_features
self.out_features= out_features
self.W= nn.Linear(in_features, out_features, bias=bias)
self.tGraphConvolution= torch.randn(2,3)
self.init()definit(self):
stdv=1./ math.sqrt(self.W.weight.size(1))
self.W.weight.data.uniform_(-stdv, stdv)defforward(self,input, adj):
support= self.W(input)# XW
output= torch.spmm(adj, support)# AXWreturn outputclassGCN(nn.Module):"""
A Two-layer GCN.
"""def__init__(self, nfeat, nhid, nclass, dropout, degree):super(GCN, self).__init__()
self.gc1= GraphConvolution(nfeat, nhid)
self.gc2= GraphConvolution(nhid, nclass)
self.dropout= dropout
self.degree= degree# self.alpha = torch.ones(self.degree, 1, requires_grad=True)
self.alpha= nn.Embedding(self.degree,1)# self.alpha.weight.data[0][0] 取出第一个元素
self.tGCN= torch.randn(2,3)defforward(self, x, adj, use_relu=True):
x= self.gc1(x, adj)if use_relu:
x= F.relu(x)
x= F.dropout(x, self.dropout, training=self.training)
x= self.gc2(x, adj)return x
model= GCN(8,8,128,0.2,2)
model= model.cuda()# 作用相同# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# model.to(device)print("over")
程序在自己的笔记本上运行,有一块cpu,有gpu,运行到model = model.cuda()
语句前,模型的所有变量都默认保存在cpu当中,model = model.cuda()
运行之后,像
self.gc1
self.W
self.alpha
都由原来的cpu设备转移到gpu了,但是对于自定义的
self.tGCN= torch.randn(2,3)
self.tGraphConvolution= torch.randn(2,3)
仍保留在cpu设备上。
不同类型的tensor相加
torch.tensor(0., device='cpu')+torch.tensor(0., device='cuda') cpu
torch.tensor(0., device='cuda')+torch.tensor(0., device='cuda') cuda
torch.tensor(0., device='cpu')+torch.tensor(0., device='cpu') cpu
设置所有的tensor为cuda类型
torch.set_default_tensor_type('torch.cuda.FloatTensor')
加上该语句即使没有model = model.cuda()
,
self.gc1
self.W
self.alpha
和
self.tGCN= torch.randn(2,3)
self.tGraphConvolution= torch.randn(2,3)
都会保存到cuda设备上