[onert/train] Attach auxiliary tensors to tensor builder

[zetwhite]

Background

backendtrain/ops/*Layer has extra(auxiliary) tensors used for backward().

For example,

ONE/runtime/onert/backend/train/ops/FullyConnectedLayer.h

Lines 61 to 65 in 60683ad

	// TODO Optimize memory
	std::unique_ptr<Tensor> _transposed_weights;
	std::unique_ptr<Tensor> _transposed_input;
	std::unique_ptr<Tensor> _transposed_back_prop_output;
	std::unique_ptr<Tensor> _act_back_prop_output;

ONE/runtime/onert/backend/train/ops/ConvolutionLayer.h

Lines 56 to 60 in 60683ad

	// TODO Consider if these tensors should be built in TensorBuilder
	std::unique_ptr<Tensor> _transposed_weights;
	std::unique_ptr<BackPropTensor> _conv_back_prop_output;
	std::unique_ptr<BackPropTensor> _act_back_prop_output;
	std::unique_ptr<GradientTensor> _transposed_grad_weights;

These tensors are allocated when KernelGenerator visits each operation.

ONE/runtime/onert/backend/train/ops/FullyConnectedLayer.cc

Lines 88 to 96 in 60683ad

	_transposed_weights = createTransposedTensor(weights);
	_transposed_weights->setBuffer(std::make_shared<basic::Allocator>(weights->total_size()));
	_transposed_input = createTransposedTensor(input);
	_transposed_input->setBuffer(std::make_shared<basic::Allocator>(input->total_size()));
	_transposed_back_prop_output = createTransposedTensor(back_prop_output);
	_transposed_back_prop_output->setBuffer(
	std::make_shared<basic::Allocator>(back_prop_output->total_size()));

What

These auxiliary tensors always hold memory after being configured.
So, Adding these tensors into TensorBuilder to use a memory planner might be helpful.

[zetwhite]

After some work on #13486,
I checked how much memory was reduced compared to the master branch.

mnist

• model : [onert] Support mnist model training #11389

• master

  • optimizer : adam
  • batchsize : 32
  • 54.2 MB (56874848 bytes)

• on draft ( Draft : Add infra to generate extra tensor #13486 )

  • optimizer : adam
  • batchsize : 32
  • 38.4 MB (40314240 bytes)

• extra tensors allocation reduces : 23187968(22.1MB) -> 6627328 (6.3 MB)

	[      ALLOC     ] allocation capacity: 11360128 # non-const 
	[      ALLOC     ] allocation capacity: 1938880  # trainable 
	[      ALLOC     ] allocation capacity: 11360000 # back-prop 
	[      ALLOC     ] allocation capacity: 1938880  # gradient
	[      ALLOC     ] allocation capacity: 3877760  #  opt variable
	[      ALLOC     ] allocation capacity: 3211264  # disposable 
	[      ALLOC     ] allocation capacity: 6627328  # extra tensors 

mobile net v2

• model : [onert] Support mobilenet_v2 model training #12325

• master

  • optimizer : adam
  • batchsize : 10
  • 1217.3 MB (1276482848 bytes)

• on draft ( Draft : Add infra to generate extra tensor #13486 )

  • optimizer : adam
  • batch : 10
  • 881.1MB ( 923913328 Bytes)

• extra tensors allocation reduces : 448920128(428.1MB) -> 96350208 (91.8MB)

	[      ALLOC     ] allocation capacity: 361362288 # non-const
	[      ALLOC     ] allocation capacity: 13951408  # trainable 
	[      ALLOC     ] allocation capacity: 361362240 # back-prop
	[      ALLOC     ] allocation capacity: 13951408  # gradient 
	[      ALLOC     ] allocation capacity: 27902816  # opt variable 
	[      ALLOC     ] allocation capacity: 49032960  # disposable
	[      ALLOC     ] allocation capacity: 96350208  # extra tensors 

/cc @ragmani

[ragmani]

After applying all PRs related to the draft #13305, the other allocation capacity will be reduces as follows:

mnist

33686912(32.1 MB) -> 25187648(24.0MB)

non-const : 11360128 -> 11341056
trainable : 1938880 -> 1938880
back-prop : 11360000 -> 6423808
gradient : 1938880 -> 1606144
optimizer variables : 3877760 -> 3877760
disposable : 3211264 -> 0

mobile net v2

827562720(789.2MB) -> 490938032(468.1MB)
827562720(789.2MB) -> 508592656(485.0MB)

non-const : 361362288 -> 361362240
trainable : 13951312 -> 13951312
back-prop : 361362240 -> 97241920
gradient : 13951312 -> 5124000
~optimizer variables : 27902608 -> 10248000~
optimizer variables : 27902608 -> 27902624
disposable : 49032960 -> 3010560

The capacity of optimizer variable is 27902624, not 10248000

[zetwhite]

I'll start to make PR based on the draft. ( #13486 )
Draft is somehow rough, I'll trim it while making a PR.

TODO

• Core

  • Introduce ExtraTensor [onert] Introduce ExtraTensor #13576
  • Rename ExtraTensor -> LayerScopeTensor [onert] Rename ExtraTensor to LayerScopeTensor #13845
  • Add LiftimeInfo to LayerScopeTensor [onert] Add lifetime enum to LayerScopeTensor #13861
  • Update ITrainableFunction [onert] Add registerLayerScopeTensors to ITrainableFuntion #13980

• Backend/train

  • Introduce ExtraTensorIndex [onert] Introduce ExtraTensorIndex #13605
  • Templatize memory planners, memory planner factory, memory manager
    • [onert] Templatize memory planners in train backend #13892
    • [onert/backend] Templatize MemoryPlannerFactory in train #13984
    • [onert/backend] Introduce LayerScopeMemoryManager #14021
  • Update TensorRegistry to have LayerScopeTensorIndex, LayerScopeTensor
    • [onert/backend] Add LayerScopeTensors into TensorRegistry #14024
  • Update TensorManager to have LayerScopeMemoryManager
    • [onert/train] Add LayerScopeManager into TensorManager #14047
  • Update TensorBuilder and TensorPlanner
    • [onert/backend] Add LayerScopeTensor interface to TensorBuilder #14198
    • [onert/backend] Add LayerScopeTensor Planning #14199
    • [onert/train] Register LayerScopeTensor to registry #14235
  • Ops
    • BianryArithmetic
    • Convolution
    • Depthwise Covolution
    • FC
    • Pool