torch.proto

syntax = "proto2";

import "caffe2/proto/caffe2.proto";

package torch;

message RecordRef {
  optional string key = 1;
}

message TensorDef {
  repeated int64 dims = 1;
  optional int64 offset = 2;
  repeated int64 strides = 3;
  // whether we compute the gradient for the parameter
  optional bool requires_grad = 4;
  optional caffe2.TensorProto.DataType data_type = 5;

  optional RecordRef data = 6;

  // device field stores the canonical device string, and it follows the
  // format below: `(cpu|cuda)[:<device-index>]`, e.g., 'cuda:0'
  optional string device = 7;

  optional bool is_quantized = 8;
  optional double scale = 9;
  optional int64 zero_point = 10;
}

message AttributeDef {
  // The mypy type of this attribute
  required string type = 1;
  required string name = 2;

  // Offset into attribute table
  required int64 id = 3;
}

message ParameterDef {
  // whether this parameter is registered as buffer or not
  optional bool is_buffer = 1;

  // the offset into the tensor table where this parameter is stored
  optional int64 tensor_id = 2;

  optional string name = 3;
}

message ModuleDef {
  repeated ModuleDef submodules = 1;

  optional RecordRef torchscript_arena = 2;

  repeated caffe2.NetDef caffe2_nets = 3;

  // because the old pickle modules may not be supported by torch_script,
  // have to stored as pickle_arena at this moment.
  optional RecordRef pickle_arena = 4;
  // should be exposed by the Class Archive, so user can save
  // module specific data which cannot be store in the graph or torch_script
  optional RecordRef cpp_arena = 5;

  // the parameters of this module
  repeated ParameterDef parameters = 6;

  // the names of inputs and outputs of the module are inferred
  // from the main method.

  optional string name = 7;

  // whether apply the optimizations to this module, only applicable to
  // script modules
  optional bool optimize = 8;

  repeated AttributeDef attributes = 9;

  // Used for retrieving module state from the pickled IValues table
  optional int64 get_state_attribute_id = 10;

  optional RecordRef torchscript_debug_arena = 11;
}

// Represents all non-module code that the model depends on.
// Right now it's just a straight list of classes, defined in dependency order
// (i.e. dependencies appear before their dependers)
message LibDef {
  optional RecordRef torchscript_arena = 1;
}

enum ProtoVersion {
  PROTO_VERSION_NEWEST = 0x0000000000000006;
}

message ModelDef {
  // numbers of fields that have been removed. Do not reuse them!
  reserved 9;
  reserved "libs";
  // for the proto version, to keep both backward and forward
  // compatibility, please bump the proto_version when we add any
  // change in the proto. runtime decides whether accept the
  // model based on the ir_version.
  optional int64 proto_version = 1;

  // main module of the model
  optional ModuleDef main_module = 2;

  // to distinguish whether exported from c2 or torch
  optional string producer_name = 3;

  // put build version here
  optional string producer_version = 4;

  // the table contains all the tensor information
  // the tensor id is defined as TensorProto.name
  repeated TensorDef tensors = 5;
}