Repository: pytorch
Issue: Well documented, safe method to deserialize model parameters from untrusted sources
Labels: feature,module: pickle,module: serialization,triaged,onnx-triaged,topic: security

Hi, I propose to add a safe way to store/load model weights (safe = can safely use weights downloaded from the internet). My vision:
- a serialization mechanism that is currently officially supported uses pickle (protocol=2) to store/load weights. Usage described [here](http://pytorch.org/docs/master/notes/serialization.html)
- pickle is unsafe 
- pytorch supports exporting to ONNX, but not importing

# Proposal 

Use ONNX format, but use only initializers field to store a model's state.

## Pros

- Safe
- Relies on already existing format with good tooling; needed protobuf-code is already in pytorch source. If needed, weights can be read from other languages
- Proposed implementation (see later) stores parameters that are met in the model state several times, loading does not break this situation (if parameters used the same tensor, they use the same tensor after loading)
- Proposed implementation was checked to normally import / export all (largest) torchvision models: inception_v3, alexnet, densenet201, resnet152, squeezenet1_1, vgg19_bn.

## Cons 

- No support for serialization of sparse tensors. AFAIK, ONNX has no support for SparseTensors so far. Hacks possible to save indices and values, but it is better to wait for official support of sparse tensors in ONNX.

## Proposed implementation

```python
import torch
from onnx import ModelProto, GraphProto, numpy_helper, load_from_string

def save_model_state(model, filename):
    graph = GraphProto()
    # exporting only once
    exported_ids = set()
    for name, tensor in model.state_dict().items():
        if id(tensor) in exported_ids:
            continue 
        else:
            exported_ids.add(id(tensor))
        try:
            numpy_value = tensor.clone().cpu().numpy()
        except:
            raise RuntimeError("Parameter {}, tensor: {} can't be dumped. \
                               Sparse tensors can't be saved".format(name, tensor))
        # cloning to avoid moving tensor from/to GPU
        initializer = numpy_helper.from_array(numpy_value, name=name)
        graph.initializer.extend([initializer])
    with open(filename, mode='wb') as f:
        f.write(ModelProto(graph=graph).SerializeToString())
        
def load_model_state(model, filename, strict=True):
    with open(filename, mode='rb') as f:
        graph_loaded = load_from_string(f.read()).graph
        
    # exporting only once
    imported_ids = set()
    own_state = model.state_dict()
    for initializer in graph_loaded.initializer:
        name = initializer.name
        if name in own_state:
            try:
                tensor = own_state[name]
                numpy_value = numpy_helper.to_array(initializer)
                tensor.copy_(tensor.new(numpy_value)[:], broadcast=False)
                imported_ids.add(id(tensor))
            except Exception:
                raise RuntimeError('While copying the parameter named {}, '
                                   'whose dimensions in the model are {} and '
                                   'whose dimensions in the checkpoint are {}.'
                                   .format(name, own_state[name].size(), numpy_value.shape))
        elif strict:
            raise KeyError('unexpected key "{}" in state_dict'.format(name))
    if strict:
        # checking that all tensors were covered
        for name, tensor in own_state.items():
            if id(tensor) not in imported_ids:
                raise KeyError('missing keys in state_dict: "{}"'.format(name))
```

Comment on implementation:
- wasn't able to use pytorch's ONNX classes, so used onnx package. Probably, worth fixing

cc @ezyang @gchanan @mruberry @BowenBao @neginraoof @houseroad @spandantiwari @lara-hdr