python torch节约内存开销的运算方法

  先是几个小示例，然后结合示例一起进行分析总结。示例如下，首先创建两个两行一列，值全为1的张量a和b，然后将b加到a上，使用python自带的id函数来获取a和b的内存地址。下面依次演示不同方法获得结果的差异：
  

1. +运算

a = torch.ones(2,1)
b = torch.ones(2,1)
before_id = id(a)
a = a + b
after_id = id(a)
before_id == after_id

  输出为：

False

  

2. +=运算

a = torch.ones(2,1)
b = torch.ones(2,1)
before_id = id(a)
a += b
after_id = id(a)
before_id == after_id

  输出为：

True

  

3. [:]运算

a = torch.ones(2,1)
b = torch.ones(2,1)
before_id = id(a)
a[:] = a + b
after_id = id(a)
before_id == after_id

  输出为：

True

  

4 总结

  从上面三个示例可以看出，+运算不是原地操作，我们将取消引用a指向的张量，转而指向新分配的内存处的张量。这可能是不可取的，原因有两个：首先，我们不想总是不必要地分配内存。 在机器学习中，我们可能有数百兆的参数，并且在一秒内多次更新所有参数。 通常情况下，我们希望原地执行这些更新。 其次，如果我们不原地更新，其他引用仍然会指向旧的内存位置， 这样我们的某些代码可能会无意中引用旧的参数。
  幸运的是，执行原地操作非常简单。 我们可以使用+=的方式进行，也可以使用切片表示法[:]将操作的结果分配给先前分配的数组，这样可以减少操作的内存开销。