Feature/reduce-decoder-mem-usage

[cathalobrien]

This change increases the memory saved from using chunking in the mapper. At the moment we use two arrays to accumulate chunks, this replaces it with a single array. At 9km this reduces peak memory usage by 6GB.

Below I have pictures of memory usage during the chunking part of the decoder at 9km

Before

Notice the zig-zag pattern. This is from the 'out1' tensor being constantly created and freed each chunk.

After

Now the zig-zag pattern is gone and peak memory usage has decreased by 6GB

[FussyDuck]

All committers have signed the CLA.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.85%. Comparing base (fd2bcf1) to head (0fb033a).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop      #84   +/-   ##
========================================
  Coverage    99.85%   99.85%           
========================================
  Files           23       23           
  Lines         1374     1374           
========================================
  Hits          1372     1372           
  Misses           2        2           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ssmmnn11]

great work. Is this from a training run or inference run?

[cathalobrien]

great work. Is this from a training run or inference run?

Inference. Havent tried in training bc this only happens when num_chunks > 1

[ssmmnn11]

Absolutely, would be just interesting to check.

[cathalobrien]

didnt see any difference during training. Makes sense since we dont use chunking

[floriankrb]

Thanks for this improvement @cathalobrien , please sign the CLA (see above...)

[cathalobrien]

Thanks for this improvement @cathalobrien , please sign the CLA (see above...)

done, thanks @floriankrb

[japols]

Looks good, nice catch :)