Memory Efficiency and Speed

[ChristophReich1996]

Hey @shikishima-TasakiLab, amazing work! Really cool you implemented the Involution as a CUDA extension. Have you tested yet how much more memory efficient and fast your implementation is than the authors' original CuPy implementation?

Cheers
Christoph

[shikishima-TasakiLab]

Thanks for your question.

I have not tested the memory efficiency and speed yet.
I believe the memory efficiency and speed of my CUDA implementation is about the same as the official CuPy implementation.

However, my OpenMP implementation has a significant improvement in memory efficiency and speed on the CPU.
Also, my implementation does not require CuPy to be installed, and it supports TorchScript and 16-bit float.

[shikishima-TasakiLab]

I experimented with the following code.

import time
import torch
from involution import Involution2d
from pytorch_memlab import MemReporter

if __name__ == '__main__':
    device = torch.device("cuda:0")
    print(f'Device: {torch.cuda.get_device_name(device)}')
    print(f'{"":-<79s}')

    inv2d = Involution2d(in_channels=4, out_channels=8).to(device)

    x: torch.Tensor = torch.rand(2, 4, 128, 128).to(device)

    for j in range(10):
        start_time = time.time()
        for i in range(1000):
            y: torch.Tensor = inv2d(x)
        elapsed_time = time.time() - start_time
        print(f'{elapsed_time} [s]')
        print(f'{"":-<79s}')

    reporter = MemReporter(inv2d)

    y: torch.Tensor = inv2d(x)
    z: torch.Tensor = y.mean()

    print('Forward')
    print(f'{"":-<79s}')
    reporter.report()

    z.backward()

    print()
    print('Backward')
    print(f'{"":-<79s}')
    reporter.report()

• Output of the official CuPy implementation. (Note: For ease of use, I have reimplemented it as a python module.)

  Device: NVIDIA GeForce RTX 2080 Ti
  -------------------------------------------------------------------------------
  1.244213581085205 [s]
  -------------------------------------------------------------------------------
  0.34833860397338867 [s]
  -------------------------------------------------------------------------------
  0.34714579582214355 [s]
  -------------------------------------------------------------------------------
  0.37246155738830566 [s]
  -------------------------------------------------------------------------------
  0.3627195358276367 [s]
  -------------------------------------------------------------------------------
  0.3523681163787842 [s]
  -------------------------------------------------------------------------------
  0.36493873596191406 [s]
  -------------------------------------------------------------------------------
  0.3746926784515381 [s]
  -------------------------------------------------------------------------------
  0.3783843517303467 [s]
  -------------------------------------------------------------------------------
  0.3817591667175293 [s]
  -------------------------------------------------------------------------------
  Forward
  -------------------------------------------------------------------------------
  /opt/conda/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py:148: UserWarning: torch.distributed.reduce_op is deprecated, please use torch.distributed.ReduceOp instead
    warnings.warn("torch.distributed.reduce_op is deprecated, please use "
  Element type                                            Size  Used MEM
  -------------------------------------------------------------------------------
  Storage on cuda:0
  sigma_mapping.0.weight                                  (8,)   512.00B
  sigma_mapping.0.bias                                    (8,)   512.00B
  initial_mapping.weight                          (8, 4, 1, 1)   512.00B
  reduce_mapping.weight                           (8, 4, 1, 1)   512.00B
  span_mapping.weight                            (49, 8, 1, 1)     2.00K
  Tensor0                                                 (8,)   512.00B
  Tensor1                                                 (8,)   512.00B
  Tensor2                                     (2, 4, 128, 128)   512.00K
  Tensor3                                                 (1,)   512.00B
  Tensor4                                     (2, 8, 128, 128)     1.00M
  Tensor5                                                 (1,)   512.00B
  -------------------------------------------------------------------------------
  Total Tensors: 393706   Used Memory: 1.51M
  The allocated memory on cuda:0: 11.63M
  Memory differs due to the matrix alignment or invisible gradient buffer tensors
  -------------------------------------------------------------------------------
  
  Backward
  -------------------------------------------------------------------------------
  Element type                                            Size  Used MEM
  -------------------------------------------------------------------------------
  Storage on cuda:0
  Tensor2                                     (2, 4, 128, 128)   512.00K
  Tensor4                                     (2, 8, 128, 128)     1.00M
  Tensor5                                                 (1,)   512.00B
  initial_mapping.weight                          (8, 4, 1, 1)   512.00B
  initial_mapping.weight.grad                     (8, 4, 1, 1)   512.00B
  reduce_mapping.weight                           (8, 4, 1, 1)   512.00B
  reduce_mapping.weight.grad                      (8, 4, 1, 1)   512.00B
  span_mapping.weight                            (49, 8, 1, 1)     2.00K
  span_mapping.weight.grad                       (49, 8, 1, 1)     2.00K
  sigma_mapping.0.weight                                  (8,)   512.00B
  sigma_mapping.0.weight.grad                             (8,)   512.00B
  sigma_mapping.0.bias                                    (8,)   512.00B
  sigma_mapping.0.bias.grad                               (8,)   512.00B
  Tensor0                                                 (8,)   512.00B
  Tensor1                                                 (8,)   512.00B
  Tensor3                                                 (1,)   512.00B
  -------------------------------------------------------------------------------
  Total Tensors: 394178   Used Memory: 1.51M
  The allocated memory on cuda:0: 1.51M
  -------------------------------------------------------------------------------

• Output of my CUDA implementation.

  Device: NVIDIA GeForce RTX 2080 Ti
  1.0515551567077637 [s]
  -------------------------------------------------------------------------------
  0.3270692825317383 [s]
  -------------------------------------------------------------------------------
  0.3310987949371338 [s]
  -------------------------------------------------------------------------------
  0.32990598678588867 [s]
  -------------------------------------------------------------------------------
  0.33088064193725586 [s]
  -------------------------------------------------------------------------------
  0.33019280433654785 [s]
  -------------------------------------------------------------------------------
  0.3247184753417969 [s]
  -------------------------------------------------------------------------------
  0.3289651870727539 [s]
  -------------------------------------------------------------------------------
  0.32881712913513184 [s]
  -------------------------------------------------------------------------------
  0.32895994186401367 [s]
  -------------------------------------------------------------------------------
  Forward
  -------------------------------------------------------------------------------
  /opt/conda/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py:148: UserWarning: torch.distributed.reduce_op is deprecated, please use torch.distributed.ReduceOp instead
    warnings.warn("torch.distributed.reduce_op is deprecated, please use "
  Element type                                            Size  Used MEM
  -------------------------------------------------------------------------------
  Storage on cuda:0
  sigma_mapping.0.weight                                  (8,)   512.00B
  sigma_mapping.0.bias                                    (8,)   512.00B
  initial_mapping.weight                          (8, 4, 1, 1)   512.00B
  reduce_mapping.weight                           (8, 4, 1, 1)   512.00B
  span_mapping.weight                            (49, 8, 1, 1)     2.00K
  Tensor0                                                 (8,)   512.00B
  Tensor1                                                 (8,)   512.00B
  Tensor2                                     (2, 4, 128, 128)   512.00K
  Tensor3                                                 (1,)   512.00B
  Tensor4                                     (2, 8, 128, 128)     1.00M
  Tensor5                                                 (1,)   512.00B
  -------------------------------------------------------------------------------
  Total Tensors: 393706   Used Memory: 1.51M
  The allocated memory on cuda:0: 11.63M
  Memory differs due to the matrix alignment or invisible gradient buffer tensors
  -------------------------------------------------------------------------------
  
  Backward
  -------------------------------------------------------------------------------
  Element type                                            Size  Used MEM
  -------------------------------------------------------------------------------
  Storage on cuda:0
  sigma_mapping.0.weight                                  (8,)   512.00B
  sigma_mapping.0.weight.grad                             (8,)   512.00B
  sigma_mapping.0.bias                                    (8,)   512.00B
  sigma_mapping.0.bias.grad                               (8,)   512.00B
  initial_mapping.weight                          (8, 4, 1, 1)   512.00B
  initial_mapping.weight.grad                     (8, 4, 1, 1)   512.00B
  reduce_mapping.weight                           (8, 4, 1, 1)   512.00B
  reduce_mapping.weight.grad                      (8, 4, 1, 1)   512.00B
  span_mapping.weight                            (49, 8, 1, 1)     2.00K
  span_mapping.weight.grad                       (49, 8, 1, 1)     2.00K
  Tensor0                                                 (8,)   512.00B
  Tensor1                                                 (8,)   512.00B
  Tensor2                                     (2, 4, 128, 128)   512.00K
  Tensor3                                                 (1,)   512.00B
  Tensor4                                     (2, 8, 128, 128)     1.00M
  Tensor5                                                 (1,)   512.00B
  -------------------------------------------------------------------------------
  Total Tensors: 394178   Used Memory: 1.51M
  The allocated memory on cuda:0: 1.51M
  -------------------------------------------------------------------------------

There is no difference in memory efficiency between them, however, my CUDA implementation runs faster than the official CuPy implementation.

[ChristophReich1996]

Hey @shikishima-TasakiLab, thanks for the results! Very cool work! Do you consider pushing it to the official repo of the authors? I guess they would be very happy about this implementation :)

[shikishima-TasakiLab]

The author's repository is mainly for experiments in the paper, so I will keep it untouched.
Instead, I reported in the author's repository that my implementation works fast.

[ChristophReich1996]

Ah, I see. But I have another quick question what IDE setup are you using for C++/CUDA? Since the C++/CUDA support of PyCharm is not really good, I'm looking for a better solution...

[shikishima-TasakiLab]

I am using Microsoft's Visual Studio Code. This IDE supports code completion and debugging of any programming language by installing extensions. I have been using "cuda-gdb" for debugging, and this IDE has recently added support for CUDA debugging as well.

P.S. When I reported my implementation to the author, he asked me to make a PR, so I ended up making one.