pytorch 设置变量的device=‘cpu’ or ‘cuda‘ |
您所在的位置：网站首页 › used和using › pytorch 设置变量的device=‘cpu’ or ‘cuda‘ |
|
背景
在电脑上运行深度学习模型，需要设置设备在cup还是GPU上运行，GPU运行速度是显著高于CPU的，但对于在不同设备上保存的变量在运行时很容易发生错误。因此需要注意 例子 from sklearn.metrics import f1_score import torch import torch.nn as nn from torch.nn import Module import torch.nn.functional as F import math import torch.optim as optim import numpy as np class GraphConvolution(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() def init(self): stdv = 1. / math.sqrt(self.W.weight.size(1)) self.W.weight.data.uniform_(-stdv, stdv) def forward(self, input, adj): support = self.W(input) # XW output = torch.spmm(adj, support) # AXW return output class GCN(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) def forward(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设备上 |
| CopyRight 2018-2019 办公设备维修网 版权所有 豫ICP备15022753号-3 |