pytorch 设置变量的device=‘cpu’ or ‘cuda‘

2022-10-09 10:07:37

背景

在电脑上运行深度学习模型,需要设置设备在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设备上

  • 作者:snake_seeker
  • 原文链接:https://blog.csdn.net/snake_seeker/article/details/116164941
    更新时间:2022-10-09 10:07:37