TrainingExtensions/tensorflow/src/python/aimet_tensorflow/common/connectedgraph.py
def _create_param_products(self, op: Op, products_dict: dict):
        """ Create products for parameters of select modules """

        def create_and_connect_product(param_name: str, product_shape: tf.TensorShape, my_op: Op,
                                       param_tensor: tf.Tensor):
            """ Create product with given name, shape, and corresponding tensor.  Connect product to my_op. """

            product = Product(my_op.name + '/' + param_name, product_shape)
            product.is_parm = True
            product.add_consumer(my_op)
            product.tensor_dict[my_op] = param_tensor
            my_op.add_input(product)
            products_dict[product.name] = product
            my_op.add_param(param_name, product)
            self._num_products_made += 1

        def create_conv2d_dense_type_params(my_op: Op):
            """ Create products for conv2d, dense, depthwise conv2d, and similar """
            tf_op = my_op.get_module()

            weight_op = WeightTensorUtils.get_read_op(tf_op)
            create_and_connect_product('kernel', weight_op.outputs[0].shape, my_op, weight_op.outputs[0])

            if not BiasUtils.is_bias_none(tf_op):
                bias_op = BiasUtils.get_bias_read_op(tf_op)
                create_and_connect_product('bias', bias_op.outputs[0].shape, my_op, bias_op.outputs[0])

        def create_batchnorm_params(my_op: Op):
            """ Create products for fusedbatchnorm """
            tf_op = my_op.get_module()

            beta_tensor = BNUtils.get_beta_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('beta', beta_tensor.shape, my_op, beta_tensor)

            gamma_tensor = BNUtils.get_gamma_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('gamma', gamma_tensor.shape, my_op, gamma_tensor)

            moving_mean_tensor = BNUtils.get_moving_mean_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('moving_mean', moving_mean_tensor.shape, my_op,
                                       moving_mean_tensor)

            moving_variance_tensor = BNUtils.get_moving_variance_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('moving_variance', moving_variance_tensor.shape, my_op,
                                       moving_variance_tensor)

        def handle_default(my_op: Op):
            """ Handler for other modules """
            logger.debug("Nothing to handle for op %s", my_op.name)

        switcher = {
            "Conv2D": create_conv2d_dense_type_params,
            "Dense": create_conv2d_dense_type_params,
            "DepthwiseConv2dNative": create_conv2d_dense_type_params,
            "BatchNorm": create_batchnorm_params,
            "FusedBatchNormV3": create_batchnorm_params
        }

        handler = switcher.get(op.type, handle_default)
        handler(op)
 Create products for parameters of select modules 
Create products for parameters of select modules
Create products for parameters of select modules
def _create_param_products(self, op: Op, products_dict: dict):
        def create_and_connect_product(param_name: str, product_shape: tf.TensorShape, my_op: Op,
                                       param_tensor: tf.Tensor):
            product = Product(my_op.name + '/' + param_name, product_shape)
            product.is_parm = True
            product.add_consumer(my_op)
            product.tensor_dict[my_op] = param_tensor
            my_op.add_input(product)
            products_dict[product.name] = product
            my_op.add_param(param_name, product)
            self._num_products_made += 1
        def create_conv2d_dense_type_params(my_op: Op):
            tf_op = my_op.get_module()
            weight_op = WeightTensorUtils.get_read_op(tf_op)
            create_and_connect_product('kernel', weight_op.outputs[0].shape, my_op, weight_op.outputs[0])
            if not BiasUtils.is_bias_none(tf_op):
                bias_op = BiasUtils.get_bias_read_op(tf_op)
                create_and_connect_product('bias', bias_op.outputs[0].shape, my_op, bias_op.outputs[0])
        def create_batchnorm_params(my_op: Op):
            tf_op = my_op.get_module()
            beta_tensor = BNUtils.get_beta_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('beta', beta_tensor.shape, my_op, beta_tensor)
            gamma_tensor = BNUtils.get_gamma_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('gamma', gamma_tensor.shape, my_op, gamma_tensor)
            moving_mean_tensor = BNUtils.get_moving_mean_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('moving_mean', moving_mean_tensor.shape, my_op,
                                       moving_mean_tensor)
            moving_variance_tensor = BNUtils.get_moving_variance_read_var_op_tensor(self._graph, tf_op)
            create_and_connect_product('moving_variance', moving_variance_tensor.shape, my_op,
                                       moving_variance_tensor)
        def handle_default(my_op: Op):
            logger.debug("Nothing to handle for op %s", my_op.name)
        switcher = {
            "Conv2D": create_conv2d_dense_type_params,
            "Dense": create_conv2d_dense_type_params,
            "DepthwiseConv2dNative": create_conv2d_dense_type_params,
            "BatchNorm": create_batchnorm_params,
            "FusedBatchNormV3": create_batchnorm_params
        }
        handler = switcher.get(op.type, handle_default)
        handler(op)
def _create_param_products ( self , op : Op , products_dict : dict ) : def create_and_connect_product ( param_name : str , product_shape : tf . TensorShape , my_op : Op , param_tensor : tf . Tensor ) : """ Create product with given name, shape, and corresponding tensor.  Connect product to my_op. """ product = Product ( my_op . name + '/' + param_name , product_shape ) product . is_parm = True product . add_consumer ( my_op ) product . tensor_dict [ my_op ] = param_tensor my_op . add_input ( product ) products_dict [ product . name ] = product my_op . add_param ( param_name , product ) self . _num_products_made += 1 def create_conv2d_dense_type_params ( my_op : Op ) : """ Create products for conv2d, dense, depthwise conv2d, and similar """ tf_op = my_op . get_module ( ) weight_op = WeightTensorUtils . get_read_op ( tf_op ) create_and_connect_product ( 'kernel' , weight_op . outputs [ 0 ] . shape , my_op , weight_op . outputs [ 0 ] ) if not BiasUtils . is_bias_none ( tf_op ) : bias_op = BiasUtils . get_bias_read_op ( tf_op ) create_and_connect_product ( 'bias' , bias_op . outputs [ 0 ] . shape , my_op , bias_op . outputs [ 0 ] ) def create_batchnorm_params ( my_op : Op ) : """ Create products for fusedbatchnorm """ tf_op = my_op . get_module ( ) beta_tensor = BNUtils . get_beta_read_var_op_tensor ( self . _graph , tf_op ) create_and_connect_product ( 'beta' , beta_tensor . shape , my_op , beta_tensor ) gamma_tensor = BNUtils . get_gamma_read_var_op_tensor ( self . _graph , tf_op ) create_and_connect_product ( 'gamma' , gamma_tensor . shape , my_op , gamma_tensor ) moving_mean_tensor = BNUtils . get_moving_mean_read_var_op_tensor ( self . _graph , tf_op ) create_and_connect_product ( 'moving_mean' , moving_mean_tensor . shape , my_op , moving_mean_tensor ) moving_variance_tensor = BNUtils . get_moving_variance_read_var_op_tensor ( self . _graph , tf_op ) create_and_connect_product ( 'moving_variance' , moving_variance_tensor . shape , my_op , moving_variance_tensor ) def handle_default ( my_op : Op ) : """ Handler for other modules """ logger . debug ( "Nothing to handle for op %s" , my_op . name ) switcher = { "Conv2D" : create_conv2d_dense_type_params , "Dense" : create_conv2d_dense_type_params , "DepthwiseConv2dNative" : create_conv2d_dense_type_params , "BatchNorm" : create_batchnorm_params , "FusedBatchNormV3" : create_batchnorm_params } handler = switcher . get ( op . type , handle_default ) handler ( op )
Create products for parameters of select modules
Create products for parameters of select modules
""" Create products for parameters of select modules """ """ Create product with given name, shape, and corresponding tensor.  Connect product to my_op. """ """ Create products for conv2d, dense, depthwise conv2d, and similar """ """ Create products for fusedbatchnorm """ """ Handler for other modules """
self
op
Op
products_dict
dict