def get_op_mem_block(self, op_type, output_shape):
if op_type == 'WinogradTransform' or op_type == 'MatMul':
buffer_shape = list(output_shape) + [1]
mem_block = MemoryBlock(
mace_pb2.GPU_IMAGE,
calculate_image_shape(OpenCLBufferType.IN_OUT_HEIGHT,
buffer_shape))
elif op_type in ['Shape',
'InferConv2dShape',
'StridedSlice',
'Stack',
'ScalarMath']:
if len(output_shape) == 1:
mem_block = MemoryBlock(mace_pb2.CPU_BUFFER,
[output_shape[0], 1])
elif len(output_shape) == 0:
mem_block = MemoryBlock(mace_pb2.CPU_BUFFER,
[1, 1])
else:
raise Exception('%s output shape dim size is not 0 or 1.' %
op_type)
else:
if len(output_shape) == 2: # only support fc/softmax
buffer_shape = [output_shape[0], 1, 1, output_shape[1]]
elif len(output_shape) == 4:
buffer_shape = output_shape
else:
raise Exception('%s output shape dim size is not 2 or 4.' %
op_type)
mem_block = MemoryBlock(
mace_pb2.GPU_IMAGE,
calculate_image_shape(OpenCLBufferType.IN_OUT_CHANNEL,
buffer_shape))
return mem_block
def get_op_mem_block(self, op_type, output_shape):
if op_type == 'WinogradTransform' or op_type == 'MatMul':
buffer_shape = list(output_shape) + [1]
mem_block = MemoryBlock(
mace_pb2.GPU_IMAGE,
calculate_image_shape(OpenCLBufferType.IN_OUT_HEIGHT,
buffer_shape))
elif op_type in [
'Shape', 'InferConv2dShape', 'StridedSlice', 'Stack',
'ScalarMath'
]:
if len(output_shape) == 1:
mem_block = MemoryBlock(mace_pb2.CPU_BUFFER,
[output_shape[0], 1])
elif len(output_shape) == 0:
mem_block = MemoryBlock(mace_pb2.CPU_BUFFER, [1, 1])
else:
raise Exception('%s output shape dim size is not 0 or 1.' %
op_type)
else:
if len(output_shape) == 2: # only support fc/softmax
buffer_shape = [output_shape[0], output_shape[1]]
elif len(output_shape) == 4:
buffer_shape = output_shape
else:
raise Exception('%s output shape dim size is not 2 or 4.' %
op_type)
mem_block = MemoryBlock(
mace_pb2.GPU_IMAGE,
calculate_image_shape(OpenCLBufferType.IN_OUT_CHANNEL,
buffer_shape))
return mem_block
def get_op_mem_block(self, op_type, output_shape):
if op_type == 'WinogradTransform' or op_type == 'MatMul':
buffer_shape = list(output_shape) + [1]
mem_block = MemoryBlock(
mace_pb2.GPU_IMAGE,
calculate_image_shape(OpenCLBufferType.IN_OUT_HEIGHT,
buffer_shape))
elif op_type == 'Shape':
mem_block = MemoryBlock(mace_pb2.CPU_BUFFER, [output_shape[0], 1])
else:
if len(output_shape) == 2: # only support fc/softmax
buffer_shape = [output_shape[0], 1, 1, output_shape[1]]
elif len(output_shape) == 4:
buffer_shape = output_shape
else:
raise Exception('output shape dim size is not 2 or 4.')
mem_block = MemoryBlock(
mace_pb2.GPU_IMAGE,
calculate_image_shape(OpenCLBufferType.IN_OUT_CHANNEL,
buffer_shape))
return mem_block