[Pytorch] torch.save原理

最新推荐文章于 2024-08-27 23:43:52 发布

原创

最新推荐文章于 2024-08-27 23:43:52 发布 · 1.1k 阅读

2 ·

CC 4.0 BY-SA版权

该博客详细解析了PyTorch中torch.save的实现原理，涉及Tensor的序列化过程，包括如何保存Storage和Tensor的数据。在序列化过程中，Tensor的Storage被转换为serialized_storages，而Tensor的其他信息通过pickle模块写入文件。在反序列化时，将分别读取这些数据重建Tensor。整个过程涉及到Python和C++的交互，以及at::StorageImpl的data()函数调用。

torch.save调用serialization.py里的_save:

def _save(obj, zip_file, pickle_module, pickle_protocol):
    serialized_storages = {}
    id_map: Dict[int, str] = {}

    def persistent_id(obj):
        # FIXME: the docs say that persistent_id should only return a string
        # but torch store returns tuples. This works only in the binary protocol
        # see
        # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects
        # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537
        if torch.is_storage(obj):
            storage_type = normalize_storage_type(type(obj))
            obj_key = id_map.setdefault(obj._cdata, str(len(id_map)))
            location = location_tag(obj)
            serialized_storages[obj_key] = obj

            return ('storage',
                    storage_type,
                    obj_key,
                    location,
                    obj.size())
        return None

    # Write the