研究で使っている保存/読出しユーティリティメソッド

研究をしていると，実験結果や実験条件をまとめたオブジェクトを保存/読出しすることが多々あります．私は特に次のようなデータを取り扱うことが多いです．

• 実験条件や実験結果をまとめた構造化データ（JSON/.json）
• 大量の要素を持つ配列（Pickle/.pkl)
• Matplotlibグラフ（PDF/.pdf）

そのため，これらを保存/読出しする共通のメソッドを作ってみました．

機能

1. いかなる形式・拡張子のデータもsave(obj,fn)で保存，load(fn)で読出し．データの形式は拡張子で判断．
2. 既にファイルが存在するのに保存しようとすると，確認のプロンプトが出る←結構重宝する

使用例

from myautils import *
from matplotlib import pyplot as plt
fig, ax = plt.subplots()
ax.plot([0, 1, 2], [0, 4, 9])
save(fig, "Experiment.Hoge_vs_Fuga.pdf")

実装

私の気分次第で機能が追加されたりするかもです．今のところ，追加したい機能は

• PyTorchモデル（.pt）を保存/読出し．
• 保存先や読出し元を空欄で呼び出すと，場所をダイアログ形式で聞いてくる．
• PyTorch TensorをJSONとして保存しようとした時は，エラーを出すのではなく，自動的に配列に変換してあげる．

myautils.py

# 保存/読出
from typing import List, Union
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from abc import ABC, abstractmethod
import json, pickle
from pathlib import Path


class ReadWriteWrapper(ABC):
    @abstractmethod
    def is_supported_obj(self, obj: any) -> bool:
        raise NotImplementedError()

    @abstractmethod
    def ext(self) -> str:
        raise NotImplementedError()

    @abstractmethod
    def save(self, obj: any, fn: str) -> None:
        raise NotImplementedError()

    @abstractmethod
    def load(self, fn: str) -> any:
        raise NotImplementedError()


class JsonReadWriteWrapper(ReadWriteWrapper):
    def is_supported_obj(self, obj: any) -> bool:
        return True

    def ext(self) -> str:
        return "json"

    def save(self, obj: any, fn: str) -> None:
        json.dump(obj, Path(fn).open("w"), ensure_ascii=False)

    def load(self, fn: str) -> any:
        return json.load(Path(fn).open("r"))


class PickleReadWriteWrapper(ReadWriteWrapper):
    def is_supported_obj(self, obj: any) -> bool:
        return True

    def ext(self) -> str:
        return "pkl"

    def save(self, obj: any, fn: str) -> None:
        pickle.dump(obj, Path(fn).open("wb"))

    def load(self, fn: str) -> any:
        return pickle.load(Path(fn).open("rb"))


class FigureReadWriteWrapper(ReadWriteWrapper):
    def is_supported_obj(self, obj: any) -> bool:
        return isinstance(obj, Figure) or isinstance(obj, Axes)

    def ext(self) -> str:
        return "pdf"

    def save(self, obj: any, fn: str) -> None:
        if isinstance(obj, Axes):
            ax: Axes = obj
            fig = ax.get_figure()
            self.save(fig, fn)
            return
        fig: Figure = obj
        fig.savefig(fn)

    def load(self, fn: str) -> any:
        raise ValueError(f"PDFを読出しすることはできません。")


_wrappers: List[ReadWriteWrapper] = [
    JsonReadWriteWrapper(),
    PickleReadWriteWrapper(),
    FigureReadWriteWrapper(),
]


def save(obj: any, fn: Union[str, Path]) -> str:
    fp = Path(fn)
    if fp.exists():
        res = input(
            f"ファイル {fp.name} は既に存在しますが、上書きしますか？(Y/その他)"
        )
        if res != "Y":
            return "保存しませんでした。"
    saver = None
    for wrapper in _wrappers:
        if fp.name.endswith(wrapper.ext()) and wrapper.is_supported_obj(obj):
            saver = wrapper
            break
    if saver is None:
        raise ValueError(
            f"ファイル {fp.name} へ保存する予定のオブジェクトについて：この拡張子はサポートされていません。"
        )
    saver.save(obj, fn)


def load(fn: Union[str,Path]) -> any:
    fp = Path(fn)
    if not fp.exists():
        ValueError(f"ファイル {fp.name} は存在しません。")
    loader = None
    for wrapper in _wrappers:
        if fp.name.endswith(wrapper.ext()):
            loader = wrapper
            break
    if loader is None:
        raise ValueError(
            f"ファイル {fp.name} から読出しする予定のオブジェクトについて：この拡張子はサポートされていません。"
        )
    return loader.load(fn)

でもね，研究のデータは

実験条件や，データサイズの小さな実験結果の情報は，手書きでノートに書く📔もしくは手作業でMarkdownや$\LaTeX$，Wordにまとめて保存するのが一番よ🌟それが結局一番上手くいくし，論文やゼミ資料を作る時の効率が最もよい．