File size: 906 Bytes
4d3c316 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | #!/usr/bin/env python
# -*- coding: utf-8 -*-
# Multi-Grained Vision Language Pre-Training: Aligning Texts with Visual Concepts (https://arxiv.org/abs/2111.08276)
# Github: https://github.com/zengyan-97/X-VLM
# Copyright (c) 2022, ByteDance Inc.
# All rights reserved.
import io
import torch
from .hdfs_io import hopen
def load(filepath: str, **kwargs):
""" load model """
if not filepath.startswith("hdfs://"):
return torch.load(filepath, **kwargs)
with hopen(filepath, "rb") as reader:
accessor = io.BytesIO(reader.read())
state_dict = torch.load(accessor, **kwargs)
del accessor
return state_dict
def save(obj, filepath: str, **kwargs):
""" save model """
if filepath.startswith("hdfs://"):
with hopen(filepath, "wb") as writer:
torch.save(obj, writer, **kwargs)
else:
torch.save(obj, filepath, **kwargs)
|