TODO: differentiation, symbolic autograd,

TODO: fusion, operators

`prim::profile` nodes are inserted on every**use** of a value by`ProfilingRecord::instrumentBlock`. Every`prim::profile` node runs a lambda that uses a captured, initial type value and the type of an incoming tensor and merges the two into`ProfiledTensorType`

`prim::profile` nodes are inserted on every**use** of a value by`ProfilingRecord::instrumentBlock`. Every`prim::profile` node runs a lambda that uses a captured, initial type value and the type of an incoming tensor and merges the two into`ProfiledTensorType`

`prim::profile` nodes are replaced with`prim::Guard` nodes by`InsertGuards`.`prim::Guard` nodes are inserted to guarantee that beyond the guard a guarded tensor will always be of the profiled shape. This guarantee will enable optimizations and codegens to generate more efficient code.

JIT attempts to reduce the number of`prim::Guard` nodes as these nodes may interefere with optimizations.

* First,`GuardElimination::moveGuardsToDefs` tries to move`prim::Guards` to their definitions, so the guards guarding the same tensor follow the definition directly or another guard on the same tensor. This step is done in

* This ordering allows us to**coalesce** (done in`GuardElimination::coalesceGuards`) multiple guards into a single one.

JIT attempts to reduce the number of`prim::Guard` nodes as these nodes may interefere with optimizations.

* First,`GuardElimination::moveGuardsToDefs` tries to move`prim::Guards` to their definitions, so the guards guarding the same tensor follow the definition directly or another guard on the same tensor. This step is done in

* This ordering allows us to**coalesce** (done in`GuardElimination::coalesceGuards`) multiple guards into a single one.

* After guards are**coaslesced** ,`GuardElimination::eliminateGuards` attempts to eliminate more guards as follows: it inspects each operation and its inputs. It checks if inputs to the operation are guarded and also if the operation produces the consistent shapes given the guarded inputs. For example, if two inputs to`add` are guaranteed to be of shape`(2, 3)`, the output shape will also always be`(2, 3)` If this property holds, JIT is allowed to remove the guard guarding operation's output.

Lastly, JIT needs to be handle cases when the assumptions about tensor shapes fail at runtime. To handle guard failures, JIT needs to be able to run the original code i.e. the code that doesn't rely on assumptions about shapes. As guards can be inserted and moved (by Optimizer) at/to arbitrary points in a computional graph, JIT needs to be able to resume execution starting from those arbitrary points onward.

`InsertBailoutNodes` builds deoptimized versions of the original computational graph, that contain the rest of computations starting from their corresponding guard failure poins and, also, captures live values needed to execute those deoptimized graphs. In other words, the pass replaces`prim::Guard` nodes with`prim::BailOut` nodes which have the`attr::Subgraph` attributes set to the deoptimized versions of the remaining computations at their corresponding`prim::Guard`s.

`InsertBailoutNodes` builds deoptimized versions of the original computational graph, that contain the rest of computations starting from their corresponding guard failure poins and, also, captures live values needed to execute those deoptimized graphs. In other words, the pass replaces`prim::Guard` nodes with`prim::BailOut` nodes which have the`attr::Subgraph` attributes set to the deoptimized versions of the remaining computations at their corresponding`prim::Guard`s.

return x0

[export.cpp](export.cpp)

[pickler.cpp](pickler.cpp)

[import.cpp](import.cpp)

[caffe2/proto/torch.proto](../../../caffe2/proto/torch.proto)

TorchScript programs are serialized with a call to`torch.jit.save()`. The resulting file (ending in`.pt` by convention) can be loaded/executed in C++ and Python.

The`.pt` file is a zip archive (which can be opened with tools such as`unzip`) and contains:

* code - the Python printed graph of a module

*`model.json` - a JSON file of a model Protobuf def (defined in[torch.proto](caffe2/proto/torch.proto))

*`tensors/` - each of the tensors of the model, with their tensor storage stored directly in a file

*`attributes.pkl` - a Python`pickle` archive of the attributes of a module

The`model.json` contains the structure information of the model. Each model must contain one main Module, and each module may contain multiple submodules, and each module contains a bunch of parameters (tensors). We serialize the metadata for each tensor inline in`model.json` (e.g., dims, strides, record name, etc).

The`code` directory contains the Python Printed`Graph`s of the main module and its submodules.

During export a list of all the tensors in a model is created. Tensors can come from either module parameters or Tensor type attributes. Metadata about each tensor is stored in`model.json` with an index into this list. The`data` field refers to the file which contains the tensor storage data. Tensors are saved by directly writing the Tensor storage to a file.

{

"tensors": [

{

"dims": [

"40",

],

"offset":"0",

"strides": [

"800",

],

"requiresGrad":true,

"dataType":"FLOAT",

"data": {

"key":"tensors/0"

},

"device":"cpu"

}

],

}

[pickler.h](pickler.h),

[pickler.cpp](pickler.cpp),

[torch/jit/_pickle.py](../../../torch/jit/_pickle.py)

[caffe2/proto/torch.proto](../../../caffe2/proto/torch.proto)

Attributes are all module properties that are not parameters or constants. Attributes are saved in a list in the order they were defined on the module. A given module may have many attributes of different types and many submodules, each with their own attributes. Attribute metadata is recorded in`model.json`:

*`type` - the full type of the attribute (in[Mypy syntax](https://mypy.readthedocs.io/en/latest/cheat_sheet_py3.html))

*`name` - the attribute's name

*`id` - the offset into the saved list of all model attributes

In`model.json`:

{

"mainModule": {

"submodules": [

{

"attributes": [

{

"type":"Dict[str, str]",

"name":"my_submodule_dictionary",

"id":0

},

{

"type":"List[Tuple[int, int]]",

"name":"my_submodule_list",

"id":1

}

]

},

],

"attributes": [

{

"type":"Dict[str, str]",

"name":"my_main_module_dictionary",

"id":2

},

{

"type":"Tensor",

"name":"my_main_module_tensor",

"id":3

}

]

},

}

Attributes of the main module and its submodules are saved to a single file in the`zip` archive of a`.pt` file named`attributes.pkl`. Attributes are stored as a Python`pickle` archive.`pickle`'s format was chosen due to:

***user friendliness** - the attributes file can be loaded in Python with`pickle`

***size limits** - formats such as Protobuf empose size limits on total message size, whereas pickle limits are on individual values (e.g. strings cannot be longer than 4 GB)

***standard format** -`pickle` is a standard Python module with a reasonably simple format. The format is a program to be consumed by a stack machine that is detailed in Python's[`pickletools.py`](https://svn.python.org/projects/python/trunk/Lib/pickletools.py)

***built-in memoization** - for shared reference types (e.g. Tensor, string, lists, dicts)

***self describing** - a separate definition file is not needed to understand the pickled data

***eager mode save** -`torch.save()` already produces a`pickle` archive, so doing the same with attributes avoids introducing yet another format

[pickler.cpp](pickler.cpp) implements a subset of the Pickle format necessary for TorchScript models.

A single file is used for the top level module and all submodules so that attributes can reference each other and share values. Unpickling`attributes.pkl` will return a tuple of values corresponding to the attributes.

All attributes are written into the`attributes.pkl` file with the exception of tensors, which store only a tensor table index (see "tensors" above). PyTorch functions defined in[torch/jit/_pickle.py](../../../torch/jit/_pickle.py) are used to mark special data types, such as this tensor table index or specialized lists. To load the`attributes.pkl` file, use the`pickle` module in Python:

import pickle

import torch

pickle.load(open("attributes.pkl","rb"))

If for some reason you don't have PyTorch installed, you can still load`attributes.pkl` with a custom[`Unpickler`](https://docs.python.org/3/library/pickle.html#pickle.Unpickler):

import pickle

classJitUnpickler(pickle.Unpickler):

deffind_class(self,module,name):

if module!='torch.jit._pickle':

raiseRuntimeError("Unknown module")

identity=lambdax: x

if name=='build_tensor_from_id':

return identity

elif name=='build_intlist':

return identity

print(JitUnpickler(open("out_dir/out/attributes.pkl","rb")).load())

Running the following snippet produces a`ScriptModule` with several attributes.

Python's`pickletools` module can be used to decode the binary blob of`attributes.pkl` into a human readable format.

import pickletools

import zipfile

import torch

from typingimport Tuple, List

classM(torch.jit.ScriptModule):

def__init__(self):

super(M,self).__init__()

self.float= torch.jit.Attribute(2.3,float)

self.tuple= torch.jit.Attribute((1,2,3,4), Tuple[int,int,int,int])

self.tensor= torch.jit.Attribute(torch.randn(2,2), torch.Tensor)

self.int_list= torch.jit.Attribute([1,2,3,4], List[int])

defforward(self):

return (self.float,self.tuple,self.tensor,self.int_list)

M().save("out.zip")

model_zip= zipfile.ZipFile("out.zip",'r')

model_zip.extractall("out_dir")

pickletools.dis(open("out_dir/out/attributes.pkl","rb"))

The output of the above commands demonstrates the concepts described earlier. Attributes are wrapped in with`2: EMPTY_LIST` and appear in the order they are defined on the module. Functions for certain special types (e.g.`List[int]`,`Tensor`) can be seen at`37: GLOBAL` and`66: GLOBAL`, followed by data specific to that type, then finally by an instruction to build the object at`65: BUILD` and`113: BUILD` respectively.

[export.cpp](export.cpp) and[import.cpp](import.cpp) handle producing the proper protobuf definitions and (de)serializing tensor data. They use[pickler.h](pickler.cpp) which implements a subset of the`pickle` stack machine.

TODO: Script Module, torch.jit.trace,__constant__ handling, weak script modules