def extract(cls, node): # update the attributes of the node block_size = node.pb.attr['block_size'].i data_format = node.pb.attr['data_format'].s.decode('utf-8') DepthToSpaceOp.update_node_stat(node, { 'block_size': block_size, 'data_format': data_format }) return cls.enabled
def extract(cls, node): # update the attributes of the node node_name = node.soft_get('name', node.id) block_size = onnx_attr(node, 'blocksize', 'i', default=None) assert block_size is not None, \ 'DepthToSpace should have "blocksize" attribute specified for node {}'.format(node_name) onnx_mode = onnx_attr(node, 'mode', 's', default=b'DCR').decode() assert onnx_mode in ['DCR', 'CRD'], 'Unrecognized mode provided for DepthToSpace node {}'.format(node_name) if onnx_mode == 'DCR': mode = 'blocks_first' else: mode = 'depth_first' DepthToSpaceOp.update_node_stat(node, {'block_size': block_size, 'mode': mode}) return cls.enabled