sony / nnabla

Neural Network Libraries
https://nnabla.org/
Apache License 2.0
2.73k stars 334 forks source link

Release unused memory in auto forward mode #1127

Closed TakuyaNarihira closed 2 years ago

TakuyaNarihira commented 2 years ago

Memory usage is optimized in auto-forward mode.

An user can save the memory by releasing nnabla.Variables which are no longer needed in auto-forward mode.

Usage examples 1

x = nn.Variable()
with auto_forward():
    y = F.identity(x) # F.identity has no grad dependencies.
    y = F.identity(y) # The rebinding of y releases the previous y and its memory .

Usage examples 2

def local_scope(x):
    h = F.identity(x) # F.identity has no grad dependencies.
    y = F.identity(h)
    return y

x = nn.Variable()
with auto_forward():
    y = local_scope(x) # After exiting local_scope, the local Variable "h" and its memory are released.

Note

Because the timing to release nnabla.Variable depends on the Python GC, the timing of memory release is not determined but is highly expected to be immediate.

Performance example

The discriminator of StyleGAN2 extracted from nnabla-examples/image-generation/stylegan2-training/ can save the GPU memory usage in auto-forward mode at the almost same level as that in static-graph mode.

Static Dynamic (after) Dynamic (before)
3168 MB 3187 MB 8143 MB