python保存numpy数据中间结果到本地文件

当我们保存一些数据的中间结果时，往往比较麻烦，这里写一个工具函数：

def get_data(mark: str = None, array: np.ndarray = None) -> np.ndarray:
    """用于记录一些numpy的数据
    传递一个mark用来标识数据含义，还有一份array的数据

    如果mark代表的数据不存在，就保存array到本地文件中
    如果mark代表的数据存在，就读取对应的本地文件并返回
    """
    assert mark is not None
    os.chdir(os.getcwd())
    hash_file = hashlib.md5(mark.encode()).hexdigest() + ".pkl"
    if os.path.exists(hash_file):
        with open(hash_file, 'rb') as file:
            array = pickle.load(file)
        print('读取本地文件')
    else:
        assert array is not None
        with open(hash_file, 'wb') as file:
            pickle.dump(array, file)
        print('保存数据到本地文件')
    return array

当我们使用mark表示一个数据时，会对mark字符串用MD5加密，得到加密后的字符串，然后以此为文件名：

• 如果文件不存在，就保存array到这个文件中
• 如果文件存在，就读取文件内容并返回

import numpy as np
import pickle
import os
import hashlib


def get_data(mark: str = None, array: np.ndarray = None) -> np.ndarray:
    """用于记录一些numpy的数据
    传递一个mark用来标识数据含义，还有一份array的数据

    如果mark代表的数据不存在，就保存array到本地文件中
    如果mark代表的数据存在，就读取对应的本地文件并返回
    """
    assert mark is not None
    os.chdir(os.getcwd())
    hash_file = hashlib.md5(mark.encode()).hexdigest() + ".pkl"
    if os.path.exists(hash_file):
        with open(hash_file, 'rb') as file:
            array = pickle.load(file)
        print('读取本地文件')
    else:
        assert array is not None
        with open(hash_file, 'wb') as file:
            pickle.dump(array, file)
        print('保存数据到本地文件')
    return array


def main():
    array = np.array(np.random.random(size=(4, 5, 6, 7)))
    hash_value = get_data(mark="数据标识", array=array)
    

if __name__ == '__main__':
    main()