forked from nudt_dsp/netrans
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import copy
|
|
from netranslib.layer.customlayer import CustomLayer
|
|
from netranslib.layer.netranslayer import IoMap
|
|
from netranslib.core.shape import Shape
|
|
from netranslib.core.dtype import DType
|
|
from netranslib.xtf import xtf as tf
|
|
from netranslib.netranslog import NetransLog as al
|
|
from netranslib.layer.layer_params import DefParam
|
|
import numpy as np
|
|
|
|
class ArgMax(CustomLayer):
|
|
|
|
op = 'argmax'
|
|
|
|
# label, description
|
|
def_input = [IoMap('in0', 'in', 'input port')]
|
|
def_output = [IoMap('out0', 'out', 'output port')]
|
|
|
|
def_param = [
|
|
DefParam('output_type', DType.int64, False),
|
|
DefParam('keepdims', False, False),
|
|
]
|
|
|
|
def setup(self, inputs, outputs):
|
|
in_shape = inputs[0].shape.dims
|
|
out_shape = copy.deepcopy(in_shape)
|
|
axis = self.params.axis
|
|
if axis < 0:
|
|
axis = len(in_shape) + axis
|
|
setattr(self.params, 'axis', axis)
|
|
if self.params.keepdims is False:
|
|
out_shape.remove(in_shape[self.params.axis])
|
|
else:
|
|
out_shape[self.params.axis] = 1
|
|
outputs[0].shape = Shape(out_shape)
|
|
|
|
def load_params_from_caffe(self, cl):
|
|
p = dict()
|
|
p['axis'] = cl.argmax_param.axis
|
|
self.set_params(p)
|
|
|
|
def load_params_from_tf(self, ruler, layer_alias, op_alias_map, tensor_data_map, anet=None):
|
|
p = dict()
|
|
p['axis'] = int(tensor_data_map['C:out0'])
|
|
self.set_params(p)
|
|
|
|
def load_params_from_onnx(self, node):
|
|
p = dict()
|
|
p['axis'] = 0
|
|
if 'axis' in node.params:
|
|
p['axis'] = node.params['axis']
|
|
if 'keepdims' in node.params and node.params['keepdims'] == 1:
|
|
p['keepdims'] = True
|
|
else:
|
|
p['keepdims'] = False
|
|
if 'select_last_index' in node.params and node.params['select_last_index'] == 1:
|
|
al.w("Netrans don't support the select_last_index = ture")
|
|
self.set_params(p)
|
|
|
|
def compute_out_tensor(self, tensor, input_tensor):
|
|
p = self.params
|
|
output_type = DType.map_to_backend_dtype(p.output_type, 'tensorflow')
|
|
out = tf.argmax(input_tensor[0], self.params.axis, output_type=output_type)
|
|
if p.keepdims is True:
|
|
out = tf.reshape(out, self.get_out_shape().dims)
|
|
return [out]
|
|
|
|
def compute_out_type(self):
|
|
p = self.params
|
|
self.get_output(0).type = p.output_type
|