解析tensor.expand()为什么不会分配新的内存而只是在存在的张量上创建新的视图
关于tensor.expand()函数的介绍见前文pytorch中的expand()和expand_as()函数。
在faster_rcnn源码的rpn模块中反复出现了view().expand().contiguous()的组合,expand()函数的功能是用来扩展张量中某维数据的尺寸,这并不难理解,但在expand()函数的解析中指出“扩展张量不会分配新的内存,只是在存在的张量上创建一个新的视图”,这句话使我有一些不理解,理解这句话的含义首先需要了解张量的存储机制。
1.张量的存储机制:
参考资料:由浅入深地了解张量
张量可以理解为一段内存的视图,多个张量可以对相同的存储进行索引,索引可以不同,但是底层内存完全相同。
张量通过使用形状(shape)/步长(stride)和存储偏移量来对相应的内存进行索引。
1)形状:表示各个维度上的元素个数,如(3, 2, 2);
2)步长:当索引在每个维度增加1时,必须跳过的内存中元素个数;
3)存储偏移量:存储中对应于张量第一个元素的index。
举例,如一段内存(1, 2, 3, 4, 5, 6),以shape=(2, 3),stride = (3, 1),index = 0来索引得到的张量为:
[[1, 2, 3],
[4, 5, 6]]
而以shape=(3, 2),stride = (1, 3),index = 0来索引同一段内存,得到的张量为:
[[1, 4],
[2, 5],
[3, 6]]
2.tensor.expand()为什么不会分配新的内存而只是在存在的张量上创建新的视图?
了解了张量的存储机制,已经可以理解了转置矩阵或者view()函数这些不改变元素个数的操作可以通过在原内存上创建新的视图实现,可是expand()扩展了张量的维度尺寸,也就是使张量的元素个数出现了增加,那么是如何通过不分配新的内存只是创建新的视图实现呢?
是通过将新的张量视图中的步长设为0实现的,如以下的例子可知,只要将stride设为0,将反复索引内存中的同一个元素,也就实现了expand()函数将张量的某个“1”维在不分配新内存情况下扩展为任意数值的更多维。
import torch
a = torch.tensor([[1], [2], [3]])
print(a.size())
c = a.expand(3, 3)
print(a)
print(c)
# 输出信息:
torch.Size([3, 1])
tensor([[1],
[2],
[3]])
tensor([[1, 1, 1],
[2, 2, 2],
[3, 3, 3]])