compiler/visq/visqlib/DumpFP32FM.py

   1 # Copyright (c) 2022 Samsung Electronics Co., Ltd. All Rights Reserved
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #    http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 # Script that dumps FM of FP32 model
  16 # NOTE This script runs on dalgona
  17
  18 import numpy as np
  19 import json
  20
  21 from pathlib import Path
  22
  23
  24 # Dump FP32 model's intermediate FM data and their names
  25 #
  26 # Before
  27 # self._dir/
  28 #
  29 # After
  30 # self._dir/
  31 #  tensors.json
  32 #  <TENSOR_ID>.npy
  33 # NOTE tensors.json has a dictionary {TENSOR_NAME -> TENSOR_ID}
  34 class DumpFP32FM:
  35     def StartAnalysis(self, args):
  36         self._dir = Path(args)
  37         self._num_data = 0
  38         # Dict {tensor_name -> tid}
  39         self._tname_to_tid = dict()
  40         self._tensor_count = 0
  41
  42     def EndNetworkExecution(self, outputs):
  43         self._num_data += 1
  44
  45     def DefaultOpPost(self, name, opcode, inputs, outputs):
  46         # Save intermediate FM into <tid>.npy
  47         data_path = self._dir / str(self._num_data)
  48         data_path.mkdir(parents=False, exist_ok=True)
  49         for output in outputs:
  50             name = output['name']
  51             data = output['data']
  52             if name in self._tname_to_tid:
  53                 tid = self._tname_to_tid[name]
  54             else:
  55                 tid = self._tensor_count
  56                 self._tname_to_tid[name] = tid
  57                 self._tensor_count += 1
  58
  59             np.save(str(data_path / str(tid)), data)
  60
  61     def EndAnalysis(self):
  62         # Save tensor name : tensor id pairs
  63         with open(self._dir / 'tensors.json', 'w') as f:
  64             json.dump(self._tname_to_tid, f, indent=2)