场景描述: 通过张量操作将一个float类型地张量转换成int32类型的张量,实际存储在内存中的比特不变。
首先要知道float在内存中的储存方法:IEEE-754格式标准,可以参考这篇博客:float数据在内存中的存储方法 比如,17.625在内存中是0 10000011 00011010000000000000000,十六进制是0x41 8D 00 00,如果当成整数来解读就是1099759616

如果是单个数字,可以用from_buffer实现:

int_value = ctypes.c_uint.from_buffer(ctypes.c_float(float_value)).value

但是如果是一个张量的话,一个一个去转换就太慢了,于是我进行搜索,找到了以下方法:

① frombuffer

属于是上面的from_buffer方法的张量版本,可以看官方文档,但是直接调用会报错:

>>> a = torch.tensor(17.625)
>>> b = torch.frombuffer(a,dtype=torch.int32)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: object does not implement Python buffer protocol.

因为tensor没有实现buffer接口,所以又去搜索,然后看到issue#19143提供的一个思路,将tensor转为numpy数组,就可以了,关于这个方法也可以参考官方文档,返回的ndarray和tensor共享同一块内存,对任意一个变量修改都会影响到对方。

>>> b = torch.frombuffer(a.numpy(),dtype=torch.int32)
>>> b
tensor([1099759616], dtype=torch.int32)

但是numpy没办法放在GPU上面。

② view

之前一直以为view只能改变张量的形状,没想到还有这个作用,view的官方文档。其实view的原理就是在同一块内存将tensor以不同的view展现出来,不存在内存拷贝行为。所以修改的时候一定要注意,base tensor也会被修改!!!,不过调用tensor.size返回的是不一样的。

命令行里面运行如下:

>>> a = torch.tensor(17.625)
>>> a
tensor(17.6250)
>>> b = a.view(torch.int32)
>>> b
tensor(1099759616, dtype=torch.int32)

如果报错TypeError: view() got an unexpected keyword argument 'dtype'那应该是pytorch版本不够,我的是1.11.0,之前用1.6.0的时候就不行。

③ memoryview

memoryview算是view的高级版,在我需要的场景中view就可以了,但是memoryview提供了更高级的功能。

Python本身支持memoryview,python官网给出了文档介绍,但是可能会有点抽象,这篇文章就写的非常的简单易懂:Python memoryview()。其实,我的理解就是,返回一个对象的内存形式(memory view这个名字也是很形象的),让我们可以直接操作这个对象所在的内存,而不是它的拷贝,也可以通过buffer protocol进行高效的byte move,所以还是很有用的。

pytorch也在尝试支持memoryview,但是好像目前这个issue#69491还没有被closed,所以应该还不能用。其实issue#69491就是基于issue#19143提出的需求而被提出的,包括issue#66382也是同样的问题,希望pytorch的开发人员搞快点吧,都从19年拖到22年了。