model-optimizer/extensions/back/PackBinaryWeights.py

   1 """
   2  Copyright (c) 2019 Intel Corporation
   3
   4  Licensed under the Apache License, Version 2.0 (the "License");
   5  you may not use this file except in compliance with the License.
   6  You may obtain a copy of the License at
   7
   8       http://www.apache.org/licenses/LICENSE-2.0
   9
  10  Unless required by applicable law or agreed to in writing, software
  11  distributed under the License is distributed on an "AS IS" BASIS,
  12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  See the License for the specific language governing permissions and
  14  limitations under the License.
  15 """
  16
  17 import logging as log
  18
  19 import networkx as nx
  20 import numpy as np
  21
  22 from mo.back.replacement import BackReplacementPattern
  23 from mo.graph.graph import Node, Graph
  24 from mo.ops.tile import Tile
  25
  26
  27 class PackBinaryWeights(BackReplacementPattern):
  28     enabled = True
  29
  30     @staticmethod
  31     def pattern():
  32         return dict(
  33             nodes=[
  34                 ('op', dict(kind='op', type='BinaryConvolution'))],
  35             edges=[]
  36         )
  37
  38     @staticmethod
  39     def replace_pattern(graph: Graph, match: dict):
  40         conv = match['op']
  41         assert len(conv.in_nodes()) == 2
  42         weights = conv.in_port(1).data.get_value().flatten()
  43         weights_rounded = np.round(weights)
  44         assert np.all(np.isclose(weights, weights_rounded))
  45         assert len(conv.in_node(1).out_nodes()) == 1
  46         weights_rounded = np.array(weights_rounded, dtype=np.int32) + 1  # -1 --> 0
  47         # Reversing element in chunks by 8 elements to pack bits correctly
  48         # First need to pad data with necessary number of element to make the length dividable by 8
  49         pad = (-len(weights_rounded))%8
  50         weights_rounded = np.array(np.concatenate((weights_rounded, np.zeros([pad]))), dtype=np.int32)
  51         assert len(weights_rounded) % 8 == 0
  52         weights_rounded = weights_rounded.reshape([len(weights_rounded)//8, 8])
  53         weights_rounded = np.flip(weights_rounded, axis=1)
  54         weights_rounded = weights_rounded.flatten()
  55         packed = np.packbits(weights_rounded)
  56         conv.in_port(1).data.set_value(packed)
  57         conv.in_node(1)['force_precision'] = 'uint8'
  58         conv['packed_weights'] = 1