Performance Pytorch：内存有效的加权和，权重沿通道共享

投入：

1） I=尺寸（N，C，X）的张量（输入）

2） W=尺寸（N，X，Y）的张量（重量）

输出：

1） O=尺寸（N，C，Y）的张量（输出）

我想计算：

I = I.view(N, C, X, 1)
W = W.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)
return O

不会产生N*C*X*Y内存开销

基本上，我想计算特征映射的加权和，其中沿通道维度的权重相同，而不会导致每个通道的内存开销

---

也许我可以用

from itertools import product

O = torch.zeros(N, C, Y)
for n, x, y in product(range(N), range(X), range(Y)):
    O[n, :, y] += I[n, :, x]*W[n, x, y]
return O

---

但是这样会比较慢（没有广播），而且我不确定为向后传递保存变量会产生多少内存开销。

您可以使用

torch.bmm

（）。只需做

torch.bmm（I，W）

要验证结果，请执行以下操作：

import torch
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)
o = torch.bmm(i,w)

# desired result code
I = i.view(N, C, X, 1)
W = w.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)

print(torch.allclose(O,o)) # should output True if outputs are same.

编辑：理想情况下，我假设使用pytorch的内部矩阵乘法是有效的。但是，您也可以使用

tracemalloc

（至少在CPU上）测量内存使用情况。有关GPU，请参阅

import torch
import tracemalloc
tracemalloc.start()
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)
o = torch.bmm(i,w)
# output is a tuple indicating current memory and peak memory
print(tracemalloc.get_traced_memory())  

您可以对其他代码执行相同的操作，并看到

bmm

实现确实是高效的

import torch
import tracemalloc
tracemalloc.start()
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)

I = i.view(N, C, X, 1)
W = w.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)
# output is a tuple indicating current memory and peak memory
print(tracemalloc.get_traced_memory())  

---

您可以使用

torch.bmm

（）。只需做

torch.bmm（I，W）

要验证结果，请执行以下操作：

import torch
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)
o = torch.bmm(i,w)

# desired result code
I = i.view(N, C, X, 1)
W = w.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)

print(torch.allclose(O,o)) # should output True if outputs are same.

编辑：理想情况下，我假设使用pytorch的内部矩阵乘法是有效的。但是，您也可以使用

tracemalloc

（至少在CPU上）测量内存使用情况。有关GPU，请参阅

import torch
import tracemalloc
tracemalloc.start()
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)
o = torch.bmm(i,w)
# output is a tuple indicating current memory and peak memory
print(tracemalloc.get_traced_memory())  

您可以对其他代码执行相同的操作，并看到

bmm

实现确实是高效的

import torch
import tracemalloc
tracemalloc.start()
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)

I = i.view(N, C, X, 1)
W = w.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)
# output is a tuple indicating current memory and peak memory
print(tracemalloc.get_traced_memory())  

---

非常感谢！但是如何检查内存消耗？我想这是最好的方法，非常感谢！但是如何检查内存消耗？我想这是最好的方法。