SimonW:


Done: https://github.com/pytorch/pytorch/issues/15968 

I submitted a patch to add those methods. Could you check the PR (linked in the issue) and let me know if they satisfy your needs?


It solves the problem. Thank you very much.
Here is a use case that others can use and feel free to add it to the test suite:
import torch
def consume_gpu_ram(n): return torch.ones((n, n)).cuda()
def consume_gpu_ram_256mb(): return consume_gpu_ram(2**13)

def b2mb(x): return int(x/2**20)
class TorchTracemalloc():

    def __enter__(self):
        self.begin = torch.cuda.memory_allocated()
        torch.cuda.reset_max_memory_allocated() # reset the peak gauge to zero
        return self

    def __exit__(self, *exc):
        self.end  = torch.cuda.memory_allocated()
        self.peak = torch.cuda.max_memory_allocated()
        self.used   = b2mb(self.end-self.begin)
        self.peaked = b2mb(self.peak-self.begin)
        print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")

# push the process' peak gauge high up and then release all the memory
# expecting 0 used / 1024 peaked
with TorchTracemalloc() as tt:
    z = [consume_gpu_ram_256mb() for i in range(4)] # 1GB
    del z
assert tt.used == 0 and tt.peaked == 1024

# allocate, allocate, release half
# expecting 256 used / 512 peaked
with TorchTracemalloc() as tt:
    # should be: 256 used, 512 peaked
    c1 = consume_gpu_ram_256mb()
    c2 = consume_gpu_ram_256mb()
    del c1
assert tt.used == 256 and tt.peaked == 512
del c2 # reset for next test

# allocate, allocate, release all
# expecting 0 used / 512 peaked
with TorchTracemalloc() as tt:
    # should be: 0 used, 512 peaked
    c1 = consume_gpu_ram_256mb()
    c2 = consume_gpu_ram_256mb()
    del c1, c2
assert tt.used == 0 and tt.peaked == 512

# allocate, don't release
# expecting 1536 used / 1536 peaked
with TorchTracemalloc() as tt:
    z = [consume_gpu_ram_256mb() for i in range(6)]
assert tt.used == 1536 and tt.peaked == 1536
del z # reset for next test

asserts already do the checks, but there is a visual check as well, output:
delta used/peak    0/1024
delta used/peak  256/ 512
delta used/peak    0/ 512
delta used/peak 1536/1536