Python TensorProto Examples

Example #1

File: utils.py Project: zxsted/caffe2

def NumpyArrayToCaffe2Tensor(arr, name):
    tensor = caffe2_pb2.TensorProto()
    tensor.data_type = caffe2_pb2.TensorProto.FLOAT
    tensor.name = name
    tensor.dims.extend(arr.shape)
    tensor.float_data.extend(list(arr.flatten().astype(float)))
    return tensor

Example #2

def NumpyArrayToCaffe2Tensor(arr, name=None):
    tensor = caffe2_pb2.TensorProto()
    tensor.dims.extend(arr.shape)
    if name:
        tensor.name = name
    if arr.dtype == np.float32:
        tensor.data_type = caffe2_pb2.TensorProto.FLOAT
        tensor.float_data.extend(list(arr.flatten().astype(float)))
    elif arr.dtype == np.float64:
        tensor.data_type = caffe2_pb2.TensorProto.DOUBLE
        tensor.double_data.extend(list(arr.flatten().astype(np.float64)))
    elif arr.dtype == np.int or arr.dtype == np.int32:
        tensor.data_type = caffe2_pb2.TensorProto.INT32
        tensor.int32_data.extend(arr.flatten().astype(np.int).tolist())
    elif arr.dtype == np.int16:
        tensor.data_type = caffe2_pb2.TensorProto.INT16
        tensor.int32_data.extend(list(arr.flatten().astype(np.int16)))  # np.int16=>pb.INT16 use int32_data
    elif arr.dtype == np.uint16:
        tensor.data_type = caffe2_pb2.TensorProto.UINT16
        tensor.int32_data.extend(list(arr.flatten().astype(np.uint16)))  # np.uint16=>pb.UNIT16 use int32_data
    elif arr.dtype == np.int8:
        tensor.data_type = caffe2_pb2.TensorProto.INT8
        tensor.int32_data.extend(list(arr.flatten().astype(np.int8)))   # np.int8=>pb.INT8 use int32_data
    elif arr.dtype == np.uint8:
        tensor.data_type = caffe2_pb2.TensorProto.UINT8
        tensor.int32_data.extend(list(arr.flatten().astype(np.uint8)))   # np.uint8=>pb.UNIT8 use int32_data
    else:
        # TODO: complete the data type: bool, float16, byte, int64, string
        raise RuntimeError(
            "Numpy data type not supported yet: " + str(arr.dtype))
    return tensor

Example #3

File: pickle_caffe_blobs.py Project: gbegkas/Detectron

 def _create_tensor(arr, shape, name):
     t = caffe2_pb2.TensorProto()
     t.name = name
     t.data_type = caffe2_pb2.TensorProto.FLOAT
     t.dims.extend(shape.dim)
     t.float_data.extend(arr)
     assert len(t.float_data) == np.prod(t.dims), 'Data size, shape mismatch'
     return t

Example #4

def NumpyArrayToCaffe2Tensor(arr, name=None):
    tensor = caffe2_pb2.TensorProto()
    tensor.dims.extend(arr.shape)
    if name:
        tensor.name = name
    if arr.dtype == np.float32:
        tensor.data_type = caffe2_pb2.TensorProto.FLOAT
        tensor.float_data.extend(list(arr.flatten().astype(float)))
    elif arr.dtype == np.float64:
        tensor.data_type = caffe2_pb2.TensorProto.DOUBLE
        tensor.double_data.extend(list(arr.flatten().astype(np.float64)))
    elif arr.dtype == np.int:
        tensor.data_type = caffe2_pb2.TensorProto.INT32
        tensor.int32_data.extend(list(arr.flatten().astype(np.int)))
    else:
        # TODO: complete the data type.
        raise RuntimeError("Numpy data type not supported yet: " +
                           str(arr.dtype))
    return tensor