将Pytorch模型从CPU转换成GPU

Posted 张乐乐章

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了将Pytorch模型从CPU转换成GPU相关的知识,希望对你有一定的参考价值。

 

1. 如何进行迁移

对模型和相应的数据进行.cuda()处理。通过这种方式,我们就可以将内存中的数据复制到GPU的显存中去。从而可以通过GPU来进行运算了。

 

1.1 判定使用GPU

下载了对应的GPU版本的Pytorch之后,要确保GPU是可以进行使用的,通过torch.cuda.is_available()的返回值来进行判断。
通过torch.cuda.device_count()可以获得能够使用的GPU数量。其他就不多赘述了。 
常常通过如下判定来写可以跑在GPU和CPU上的通用模型:

1 if torch.cuda.is_available():
2     ten1 = ten1.cuda()
3     MyModel = MyModel.cuda() 

 

 

2. 对应数据的迁移

2.1 将Tensor迁移到显存中去

不论是什么类型的Tensor(FloatTensor或者是LongTensor等等),一律直接使用方法.cuda()即可。 
例如:

1 ten1 = torch.FloatTensor(2)
2 
3 ten1_cuda = ten1.cuda()

 

如果要将显存中的数据复制到内存中,则对cuda数据类型使用.cpu()方法即可。

2.2 将Variable迁移到显存中去

在模型中,我们最常使用的是Variable这个容器来装载使用数据。主要是由于Variable可以进行反向传播来进行自动求导。 
同样地,要将Variable迁移到显存中,同样只需要使用.cuda()即可实现。

这里有一个小疑问,对Variable直接使用.cuda和对Tensor进行.cuda然后再放置到Variable中结果是否一致呢。答案是肯定的。

ten1 = torch.FloatTensor(2)
>>>  6.1101e+24
     4.5659e-41
    [torch.FloatTensor of size 2]

ten1_cuda = ten1.cuda()
>>>>  6.1101e+24
      4.5659e-41
    [torch.cuda.FloatTensor of size 2 (GPU 0)]

V1_cpu = autograd.Variable(ten1)
>>>> Variable containing:
     6.1101e+24
     4.5659e-41
    [torch.FloatTensor of size 2]

V2 = autograd.Variable(ten1_cuda)
>>>> Variable containing:
     6.1101e+24
     4.5659e-41
    [torch.cuda.FloatTensor of size 2 (GPU 0)]

V1 = V1_cpu.cuda()
>>>> Variable containing:
     6.1101e+24
     4.5659e-41
    [torch.cuda.FloatTensor of size 2 (GPU 0)]

 

最终我们能发现他们都能够达到相同的目的,但是他们完全一样了吗?我们使用V1 is V2发现,结果是否定的。

对于V1,我们是直接对Variable进行操作的,这样子V1的.grad_fn中会记录下创建的方式。因此这二者并不是完全相同的。

2.3 数据迁移小结

.cuda()操作默认使用GPU 0也就是第一张显卡来进行操作。当我们想要存储在其他显卡中时可以使用.cuda(<显卡号数>)来将数据存储在指定的显卡中。还有很多种方式,具体参考官方文档。

 

3. 模型迁移

模型的迁移这里指的是torch.nn下面的一些网络模型以及自己创建的模型迁移到GPU上去。

上面讲了使用.cuda()即可将数据从内存中移植到显存中去。 
对于模型来说,也是同样的方式,我们使用.cuda来将网络放到显存上去。

3.1 torch.nn下的基本模型迁移

我们很惊奇地发现对于模型来说,不像数据那样使用了.cuda()之后会改变其的数据类型。模型看起来没有任何的变化。 
但是他真的没有改变吗。 
我们将data1投入linear_cuda中去可以发现,系统会报错,而将.cuda之后的data2投入linear_cuda才能正常工作。并且输出的也是具有cuda的数据类型。

那是怎么一回事呢? 
这是因为这些所谓的模型,其实也就是对输入参数做了一些基本的矩阵运算。所以我们对模型.cuda()实际上也相当于将模型使用到的参数存储到了显存上去。

对于上面的例子,我们可以通过观察参数来发现区别所在。

linear.weight
>>>> Parameter containing:
    -0.6847  0.2149
    -0.5473  0.6863
    [torch.FloatTensor of size 2x2]

linear_cuda.weight
>>>> Parameter containing:
    -0.6847  0.2149
    -0.5473  0.6863
    [torch.cuda.FloatTensor of size 2x2 (GPU 0)]

 

3.2 自己模型的迁移

对于自己创建的模型类,由于继承了torch.nn.Module,则可同样使用.cuda()来将模型中用到的所有参数都存储到显存中去。

这里笔者曾经有一个疑问:当我们对模型存储到显存中去之后,那么这个模型中的方法后面所创建出来的Tensor是不是都会默认变成cuda的数据类型。答案是否定的。具体操作留给读者自己去实现。

3.3 模型小结

对于模型而言,我们可以将其看做是一种类似于Variable的容器。我们对它进行.cuda()处理,是将其中的参数放到显存上去(因为实际使用的时候也是通过这些参数做运算)。

 

https://blog.csdn.net/qq_28444159/article/details/78781201

以上是关于将Pytorch模型从CPU转换成GPU的主要内容,如果未能解决你的问题,请参考以下文章

PyTorch在GPU上训练模型

python常用代码

如何将经过 gpu 训练的模型加载到 cpu 中?

直接从 TensorFlow 访问 PyTorch GPU 矩阵

使用 cpu 与 gpu 进行训练的 pytorch 模型精度之间的巨大差异

Pytorch 模型在 CPU 和 GPU 上都内存不足,无法弄清楚我做错了啥