def ToProto()

in lingvo/core/hyperparams.py [0:0]

• 60 lines of code
• 27 McCabe index (conditional complexity)

  def ToProto(self) -> hyperparams_pb2.Hyperparam:
    """Writes to a Hyperparams proto.

    Serializes the Hyperparams into a proto that can be then written to disk or
    sent over the network. Note that serialization is not guaranteed to be
    unique or stable (this is a feature of protos themselves, not this code), so
    using it for fingerprinting for example may not be appropriate. Refer to the
    ToText() method for a serialization approach that Lingvo controls.

    Returns:
      The serialized params as a Hyperparams proto.
    """

    def _ToParamValue(key: str, val: Any) -> hyperparams_pb2.HyperparamValue:
      """Serializes to HyperparamValue proto."""
      param_pb = hyperparams_pb2.HyperparamValue()
      if isinstance(val, Params):
        param_pb.param_val.CopyFrom(_ToParam(val, prefix=key))
      elif isinstance(val, list) or isinstance(val, range):
        # The range function is serialized by explicitly calling it.
        param_pb.list_val.items.extend(
            [_ToParamValue(f'{key}[{i}]', v) for i, v in enumerate(val)])
      elif isinstance(val, tuple):
        param_pb.tuple_val.items.extend(
            [_ToParamValue(f'{key}[{i}]', v) for i, v in enumerate(val)])
      elif dataclasses.is_dataclass(val) or _IsNamedTuple(val):
        val_cls = type(val)
        items = val.__dict__.items() if dataclasses.is_dataclass(
            val) else val._asdict().items()
        param_pb.named_tuple_val.type = inspect.getmodule(
            val_cls).__name__ + '/' + val_cls.__name__
        param_pb.named_tuple_val.items.extend(
            [_ToParamValue(f'{key}[{k}]', v) for k, v in items])
      # Only dicts where all keys are str can be stored as dict_val.
      elif isinstance(val, dict) and all(isinstance(k, str) for k in val):
        param_pb.dict_val.SetInParent()
        for k, v in val.items():
          param_pb.dict_val.items[k].CopyFrom(_ToParamValue(f'{key}[{k}]', v))
      elif isinstance(val, type):
        param_pb.type_val = inspect.getmodule(val).__name__ + '/' + val.__name__
      elif isinstance(val, tf.DType):
        param_pb.dtype_val = val.name
      elif isinstance(val, str):
        param_pb.string_val = val
      elif isinstance(val, bool):
        param_pb.bool_val = val
      elif isinstance(val, int):
        param_pb.int_val = val
      elif isinstance(val, float):
        param_pb.float_val = val
      elif isinstance(val, enum.Enum):
        enum_cls = type(val)
        param_pb.enum_val.type = inspect.getmodule(
            enum_cls).__name__ + '/' + enum_cls.__name__
        param_pb.enum_val.name = val.name
      elif isinstance(val, message.Message):
        proto_cls = type(val)
        param_pb.proto_val.type = inspect.getmodule(
            proto_cls).__name__ + '/' + proto_cls.__name__
        param_pb.proto_val.val = val.SerializeToString()
      elif symbolic.IsExpr(val):
        param_pb.symbolic_val = pickle.dumps(val)
      elif val is None:
        # We represent a NoneType by the absence of any of the oneof.
        pass
      else:
        param_pb.string_repr_val = repr(val)
      return param_pb

    def _ToParam(val: 'Params', prefix: str = '') -> hyperparams_pb2.Hyperparam:
      """Serializes to Hyperparam proto."""
      param_pb = hyperparams_pb2.Hyperparam()
      if prefix:
        prefix += '.'
      for k, v in val.IterParams():
        param_pb.items[k].CopyFrom(_ToParamValue(prefix + k, v))
      return param_pb

    return _ToParam(self)