PyTorch 随机数生成占用 CPU 过高的解决方法

yipeiwu_com5年前Python基础

PyTorch 随机数生成占用 CPU 过高的问题

今天在使用 pytorch 的过程中,发现 CPU 占用率过高。经过检查,发现是因为先在 CPU 中生成了随机数,然后再调用.to(device)传到 GPU,这样导致效率变得很低,并且CPU 和 GPU 都被消耗。

查阅PyTorch文档后发现,torch.randn(shape, out)可以直接在GPU中生成随机数,只要shape是tensor.cuda.Tensor类型即可。这样,就可以避免在 CPU 中生成过大的矩阵,而 shape 变量是很小的。

因此,下面的代码就可以进行这种操作了。

noise = torch.cuda.FloatTensor(shape) if torch.cuda.is_available() else torch.FloatTensor(shape)
torch.randn(shape, out=noise)

以上这篇PyTorch 随机数生成占用 CPU 过高的解决方法就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

Python学习小技巧之列表项的拼接

本文介绍的是关于Python实现列表项拼接的一个小技巧,分享出来供大家参考学习,下面来看看详细的介绍: 典型代码: data_list = ['a', 'b', 'c', 'd',...

如何基于python实现归一化处理

如何基于python实现归一化处理

这篇文章主要介绍了如何基于python实现归一化处理,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下    ...

Python中的作用域规则详解

Python是静态作用域语言,尽管它自身是一个动态语言。也就是说,在Python中变量的作用域是由它在源代码中的位置决定的,这与C有些相似,但是Python与C在作用域方面的差异还是非常...

pyqt5利用pyqtDesigner实现登录界面

本文实例为大家分享了pyqt5利用pyqtDesigner实现登录界面的具体代码,供大家参考,具体内容如下 为便于操作 界面和逻辑分离 逻辑类: import sys import...

对Python之gzip文件读写的方法详解

gzip文件读写的时候需要用到Python的gzip模块。 具体使用如下: # -*- coding: utf-8 -*- import gzip # 写文件 f_out = gz...