| # @package utils |
| # Module caffe2.python.utils |
| |
| |
| |
| |
| |
| from caffe2.proto import caffe2_pb2 |
| from future.utils import viewitems |
| from google.protobuf.message import DecodeError, Message |
| from google.protobuf import text_format |
| |
| import sys |
| import collections |
| import copy |
| import functools |
| import numpy as np |
| from six import integer_types, binary_type, text_type, string_types |
| |
| OPTIMIZER_ITERATION_NAME = "optimizer_iteration" |
| ITERATION_MUTEX_NAME = "iteration_mutex" |
| |
| |
| def OpAlmostEqual(op_a, op_b, ignore_fields=None): |
| ''' |
| Two ops are identical except for each field in the `ignore_fields`. |
| ''' |
| ignore_fields = ignore_fields or [] |
| if not isinstance(ignore_fields, list): |
| ignore_fields = [ignore_fields] |
| |
| assert all(isinstance(f, text_type) for f in ignore_fields), ( |
| 'Expect each field is text type, but got {}'.format(ignore_fields)) |
| |
| def clean_op(op): |
| op = copy.deepcopy(op) |
| for field in ignore_fields: |
| if op.HasField(field): |
| op.ClearField(field) |
| return op |
| |
| op_a = clean_op(op_a) |
| op_b = clean_op(op_b) |
| return op_a == op_b or str(op_a) == str(op_b) |
| |
| |
| def CaffeBlobToNumpyArray(blob): |
| if (blob.num != 0): |
| # old style caffe blob. |
| return (np.asarray(blob.data, dtype=np.float32) |
| .reshape(blob.num, blob.channels, blob.height, blob.width)) |
| else: |
| # new style caffe blob. |
| return (np.asarray(blob.data, dtype=np.float32) |
| .reshape(blob.shape.dim)) |
| |
| |
| def Caffe2TensorToNumpyArray(tensor): |
| if tensor.data_type == caffe2_pb2.TensorProto.FLOAT: |
| return np.asarray( |
| tensor.float_data, dtype=np.float32).reshape(tensor.dims) |
| elif tensor.data_type == caffe2_pb2.TensorProto.DOUBLE: |
| return np.asarray( |
| tensor.double_data, dtype=np.float64).reshape(tensor.dims) |
| elif tensor.data_type == caffe2_pb2.TensorProto.INT64: |
| return np.asarray( |
| tensor.int64_data, dtype=np.int64).reshape(tensor.dims) |
| elif tensor.data_type == caffe2_pb2.TensorProto.INT32: |
| return np.asarray( |
| tensor.int32_data, dtype=np.int).reshape(tensor.dims) # pb.INT32=>np.int use int32_data |
| elif tensor.data_type == caffe2_pb2.TensorProto.INT16: |
| return np.asarray( |
| tensor.int32_data, dtype=np.int16).reshape(tensor.dims) # pb.INT16=>np.int16 use int32_data |
| elif tensor.data_type == caffe2_pb2.TensorProto.UINT16: |
| return np.asarray( |
| tensor.int32_data, dtype=np.uint16).reshape(tensor.dims) # pb.UINT16=>np.uint16 use int32_data |
| elif tensor.data_type == caffe2_pb2.TensorProto.INT8: |
| return np.asarray( |
| tensor.int32_data, dtype=np.int8).reshape(tensor.dims) # pb.INT8=>np.int8 use int32_data |
| elif tensor.data_type == caffe2_pb2.TensorProto.UINT8: |
| return np.asarray( |
| tensor.int32_data, dtype=np.uint8).reshape(tensor.dims) # pb.UINT8=>np.uint8 use int32_data |
| else: |
| # TODO: complete the data type: bool, float16, byte, int64, string |
| raise RuntimeError( |
| "Tensor data type not supported yet: " + str(tensor.data_type)) |
| |
| |
| 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.int64: |
| tensor.data_type = caffe2_pb2.TensorProto.INT64 |
| tensor.int64_data.extend(list(arr.flatten().astype(np.int64))) |
| 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, string |
| raise RuntimeError( |
| "Numpy data type not supported yet: " + str(arr.dtype)) |
| return tensor |
| |
| |
| def MakeArgument(key, value): |
| """Makes an argument based on the value type.""" |
| argument = caffe2_pb2.Argument() |
| argument.name = key |
| iterable = isinstance(value, collections.abc.Iterable) |
| |
| # Fast tracking common use case where a float32 array of tensor parameters |
| # needs to be serialized. The entire array is guaranteed to have the same |
| # dtype, so no per-element checking necessary and no need to convert each |
| # element separately. |
| if isinstance(value, np.ndarray) and value.dtype.type is np.float32: |
| argument.floats.extend(value.flatten().tolist()) |
| return argument |
| |
| if isinstance(value, np.ndarray): |
| value = value.flatten().tolist() |
| elif isinstance(value, np.generic): |
| # convert numpy scalar to native python type |
| value = np.asscalar(value) |
| |
| if type(value) is float: |
| argument.f = value |
| elif type(value) in integer_types or type(value) is bool: |
| # We make a relaxation that a boolean variable will also be stored as |
| # int. |
| argument.i = value |
| elif isinstance(value, binary_type): |
| argument.s = value |
| elif isinstance(value, text_type): |
| argument.s = value.encode('utf-8') |
| elif isinstance(value, caffe2_pb2.NetDef): |
| argument.n.CopyFrom(value) |
| elif isinstance(value, Message): |
| argument.s = value.SerializeToString() |
| elif iterable and all(type(v) in [float, np.float_] for v in value): |
| argument.floats.extend( |
| v.item() if type(v) is np.float_ else v for v in value |
| ) |
| elif iterable and all( |
| type(v) in integer_types or type(v) in [bool, np.int_] for v in value |
| ): |
| argument.ints.extend( |
| v.item() if type(v) is np.int_ else v for v in value |
| ) |
| elif iterable and all( |
| isinstance(v, binary_type) or isinstance(v, text_type) for v in value |
| ): |
| argument.strings.extend( |
| v.encode('utf-8') if isinstance(v, text_type) else v |
| for v in value |
| ) |
| elif iterable and all(isinstance(v, caffe2_pb2.NetDef) for v in value): |
| argument.nets.extend(value) |
| elif iterable and all(isinstance(v, Message) for v in value): |
| argument.strings.extend(v.SerializeToString() for v in value) |
| else: |
| if iterable: |
| raise ValueError( |
| "Unknown iterable argument type: key={} value={}, value " |
| "type={}[{}]".format( |
| key, value, type(value), set(type(v) for v in value) |
| ) |
| ) |
| else: |
| raise ValueError( |
| "Unknown argument type: key={} value={}, value type={}".format( |
| key, value, type(value) |
| ) |
| ) |
| return argument |
| |
| |
| def TryReadProtoWithClass(cls, s): |
| """Reads a protobuffer with the given proto class. |
| |
| Inputs: |
| cls: a protobuffer class. |
| s: a string of either binary or text protobuffer content. |
| |
| Outputs: |
| proto: the protobuffer of cls |
| |
| Throws: |
| google.protobuf.message.DecodeError: if we cannot decode the message. |
| """ |
| obj = cls() |
| try: |
| text_format.Parse(s, obj) |
| return obj |
| except (text_format.ParseError, UnicodeDecodeError): |
| obj.ParseFromString(s) |
| return obj |
| |
| |
| def GetContentFromProto(obj, function_map): |
| """Gets a specific field from a protocol buffer that matches the given class |
| """ |
| for cls, func in viewitems(function_map): |
| if type(obj) is cls: |
| return func(obj) |
| |
| |
| def GetContentFromProtoString(s, function_map): |
| for cls, func in viewitems(function_map): |
| try: |
| obj = TryReadProtoWithClass(cls, s) |
| return func(obj) |
| except DecodeError: |
| continue |
| else: |
| raise DecodeError("Cannot find a fit protobuffer class.") |
| |
| |
| def ConvertProtoToBinary(proto_class, filename, out_filename): |
| """Convert a text file of the given protobuf class to binary.""" |
| with open(filename) as f: |
| proto = TryReadProtoWithClass(proto_class, f.read()) |
| with open(out_filename, 'w') as fid: |
| fid.write(proto.SerializeToString()) |
| |
| |
| def GetGPUMemoryUsageStats(): |
| """Get GPU memory usage stats from CUDAContext/HIPContext. This requires flag |
| --caffe2_gpu_memory_tracking to be enabled""" |
| from caffe2.python import workspace, core |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "GetGPUMemoryUsage", |
| [], |
| ["____mem____"], |
| device_option=core.DeviceOption(workspace.GpuDeviceType, 0), |
| ), |
| ) |
| b = workspace.FetchBlob("____mem____") |
| return { |
| 'total_by_gpu': b[0, :], |
| 'max_by_gpu': b[1, :], |
| 'total': np.sum(b[0, :]), |
| 'max_total': np.sum(b[1, :]) |
| } |
| |
| |
| def ResetBlobs(blobs): |
| from caffe2.python import workspace, core |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "Free", |
| list(blobs), |
| list(blobs), |
| device_option=core.DeviceOption(caffe2_pb2.CPU), |
| ), |
| ) |
| |
| |
| class DebugMode(object): |
| ''' |
| This class allows to drop you into an interactive debugger |
| if there is an unhandled exception in your python script |
| |
| Example of usage: |
| |
| def main(): |
| # your code here |
| pass |
| |
| if __name__ == '__main__': |
| from caffe2.python.utils import DebugMode |
| DebugMode.run(main) |
| ''' |
| |
| @classmethod |
| def run(cls, func): |
| try: |
| return func() |
| except KeyboardInterrupt: |
| raise |
| except Exception: |
| import pdb |
| |
| print( |
| 'Entering interactive debugger. Type "bt" to print ' |
| 'the full stacktrace. Type "help" to see command listing.') |
| print(sys.exc_info()[1]) |
| print |
| |
| pdb.post_mortem() |
| sys.exit(1) |
| raise |
| |
| |
| def raiseIfNotEqual(a, b, msg): |
| if a != b: |
| raise Exception("{}. {} != {}".format(msg, a, b)) |
| |
| |
| def debug(f): |
| ''' |
| Use this method to decorate your function with DebugMode's functionality |
| |
| Example: |
| |
| @debug |
| def test_foo(self): |
| raise Exception("Bar") |
| |
| ''' |
| |
| @functools.wraps(f) |
| def wrapper(*args, **kwargs): |
| def func(): |
| return f(*args, **kwargs) |
| return DebugMode.run(func) |
| |
| return wrapper |
| |
| |
| def BuildUniqueMutexIter( |
| init_net, |
| net, |
| iter=None, |
| iter_mutex=None, |
| iter_val=0 |
| ): |
| ''' |
| Often, a mutex guarded iteration counter is needed. This function creates a |
| mutex iter in the net uniquely (if the iter already existing, it does |
| nothing) |
| |
| This function returns the iter blob |
| ''' |
| iter = iter if iter is not None else OPTIMIZER_ITERATION_NAME |
| iter_mutex = iter_mutex if iter_mutex is not None else ITERATION_MUTEX_NAME |
| from caffe2.python import core |
| if not init_net.BlobIsDefined(iter): |
| # Add training operators. |
| with core.DeviceScope( |
| core.DeviceOption(caffe2_pb2.CPU, |
| extra_info=["device_type_override:cpu"]) |
| ): |
| iteration = init_net.ConstantFill( |
| [], |
| iter, |
| shape=[1], |
| value=iter_val, |
| dtype=core.DataType.INT64, |
| ) |
| iter_mutex = init_net.CreateMutex([], [iter_mutex]) |
| net.AtomicIter([iter_mutex, iteration], [iteration]) |
| else: |
| iteration = init_net.GetBlobRef(iter) |
| return iteration |
| |
| |
| def EnumClassKeyVals(cls): |
| # cls can only be derived from object |
| assert type(cls) == type |
| # Enum attribute keys are all capitalized and values are strings |
| enum = {} |
| for k in dir(cls): |
| if k == k.upper(): |
| v = getattr(cls, k) |
| if isinstance(v, string_types): |
| assert v not in enum.values(), ( |
| "Failed to resolve {} as Enum: " |
| "duplicate entries {}={}, {}={}".format( |
| cls, k, v, [key for key in enum if enum[key] == v][0], v |
| ) |
| ) |
| enum[k] = v |
| return enum |
| |
| |
| def ArgsToDict(args): |
| """ |
| Convert a list of arguments to a name, value dictionary. Assumes that |
| each argument has a name. Otherwise, the argument is skipped. |
| """ |
| ans = {} |
| for arg in args: |
| if not arg.HasField("name"): |
| continue |
| for d in arg.DESCRIPTOR.fields: |
| if d.name == "name": |
| continue |
| if d.label == d.LABEL_OPTIONAL and arg.HasField(d.name): |
| ans[arg.name] = getattr(arg, d.name) |
| break |
| elif d.label == d.LABEL_REPEATED: |
| list_ = getattr(arg, d.name) |
| if len(list_) > 0: |
| ans[arg.name] = list_ |
| break |
| else: |
| ans[arg.name] = None |
| return ans |
| |
| |
| def NHWC2NCHW(tensor): |
| assert tensor.ndim >= 1 |
| return tensor.transpose((0, tensor.ndim - 1) + tuple(range(1, tensor.ndim - 1))) |
| |
| |
| def NCHW2NHWC(tensor): |
| assert tensor.ndim >= 2 |
| return tensor.transpose((0,) + tuple(range(2, tensor.ndim)) + (1,)) |
