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设备上