Stonesjtu / pytorch_memlab

Profiling and inspecting memory in pytorch
MIT License
1.01k stars 37 forks source link

Naming Tensors within MemReporter #45

Closed m-lyon closed 2 years ago

m-lyon commented 2 years ago

I'm developing a custom network layer and it subsequently has many unnamed Tensors within the MemReporter output. below is a snippet example

pcconv5_1.weight_net.layers.4.weight                 (1, 64)   512.00B
pcconv5_1.weight_net.layers.4.bias                      (1,)   512.00B
pcconv5_1.feature_layer.weight                       (1, 32)   512.00B
Tensor0                                   (1, 10, 10, 10, 8)    31.50K
Tensor1                                            (1, 8, 3)   512.00B
Tensor2                                  (1, 10, 10, 10, 12)    47.00K
Tensor3                                           (1, 12, 3)   512.00B
Tensor4                                            (1, 8, 8)   512.00B
Tensor5                                 (1, 8, 1, 1, 1, 8, 5)     1.50K
Tensor6                                 (1, 8, 1, 1, 1, 8, 32)     8.00K
Tensor7                                 (1, 8, 1, 1, 1, 8, 32)     8.00K
Tensor8                                 (1, 8, 1, 1, 1, 8, 64)    16.00K
Tensor9                                 (1, 8, 1, 1, 1, 8, 64)    16.00K
Tensor10(->Tensor4)                     (1, 8, 1, 1, 1, 8, 1)     0.00B

Is there any way to name or label these tensors (Tensors 0 through 10 for example) so that I can more easily determine which operation specifically creates them?

Stonesjtu commented 2 years ago

I'm afraid not because PyTorch does not name any tensor.

Probably we can manage a tensor -> name mapping in this tool, but it requires considerable amount of code changes.

I would recommend you to report the memory usage at multiple critical locations so you can have a deeper understand the memory system underneath.

a = torch.Tensor(5,5,5)
reporter.report()
b = a + 5
reporter.report()
c = b * 5
reporter.report()