【Python高级编程】Pickle实现AI算法训练的权重数据的保存

任务描述

在这里插入图片描述

代码实现

import pickle
import time
import os
import numpy as np


# 模拟耗时的权重计算过程
def calculate_weights():
    print("开始计算权重...")
    time.sleep(5)  # 模拟耗时操作，暂停5秒以模拟计算过程
    weights = np.random.rand(10, 10)  # 随机生成一个10x10的权重矩阵
    print("权重计算完成.")
    return weights


# 保存权重和epoch到文件
def save_weights(weights, epoch, filename='weights.pkl'):
    data = {'weights': weights, 'epoch': epoch}  # 将权重和epoch打包成字典
    with open(filename, 'wb') as f:  # 以二进制写模式打开文件
        pickle.dump(data, f)  # 使用pickle将数据序列化并保存到文件
    print(f"权重和epoch已保存到{filename}.")


# 从文件加载权重和epoch
def load_weights(filename='weights.pkl'):
    with open(filename, 'rb') as f:  # 以二进制读模式打开文件
        data = pickle.load(f)  # 使用pickle从文件中反序列化数据
    print(f"权重和epoch已从{filename}加载.")
    return data['weights'], data['epoch']  # 返回权重和epoch数据


# 主程序
def main():
    weights_file = 'weights.pkl'  # 权重文件名
    total_epochs = 100  # 假设我们总共需要训练100个epochs

    # 如果权重文件存在，则加载权重和epoch
    if os.path.exists(weights_file):
        weights, start_epoch = load_weights(weights_file)  # 从文件加载权重和训练的epoch数
    else:
        # 否则，从第一个epoch开始，并计算权重
        weights = calculate_weights()  # 计算权重
        start_epoch = 0  # 从第0个epoch开始

    # 继续训练剩余的epochs
    for epoch in range(start_epoch, total_epochs):
        print(f"开始训练epoch {epoch}...")
        # 这里进行实际的训练代码...
        time.sleep(1)  # 模拟训练过程，暂停1秒以模拟每个epoch的训练时间
        print(f"完成训练epoch {epoch}.")

        # 每个epoch结束后保存权重和epoch信息
        save_weights(weights, epoch, weights_file)  # 保存当前的权重和epoch信息

if __name__ == '__main__':
    main()

代码解释

导入必要的模块：

import pickle
import time
import os
import numpy as np

定义权重计算函数：

def calculate_weights():
    print("开始计算权重...")
    time.sleep(5)  # 模拟耗时操作，暂停5秒
    weights = np.random.rand(10, 10)  # 生成一个10x10的随机权重矩阵
    print("权重计算完成.")
    return weights

定义保存权重和epoch到文件的函数：

def save_weights(weights, epoch, filename='weights.pkl'):
    data = {'weights': weights, 'epoch': epoch}  # 将权重和epoch打包成字典
    with open(filename, 'wb') as f:  # 以二进制写模式打开文件
        pickle.dump(data, f)  # 使用pickle将字典序列化并写入文件
    print(f"权重和epoch已保存到{filename}.")

定义从文件加载权重和epoch的函数：

def load_weights(filename='weights.pkl'):
    with open(filename, 'rb') as f:  # 以二进制读模式打开文件
        data = pickle.load(f)  # 使用pickle从文件中反序列化数据
    print(f"权重和epoch已从{filename}加载.")
    return data['weights'], data['epoch']  # 返回权重和epoch数据

主程序逻辑：

def main():
    weights_file = 'weights.pkl'  # 定义权重文件的名称
    total_epochs = 100  # 假设我们需要训练100个epochs

    # 检查权重文件是否存在
    if os.path.exists(weights_file):
        weights, start_epoch = load_weights(weights_file)  # 如果存在，加载权重和epoch
    else:
        weights = calculate_weights()  # 如果不存在，计算权重
        start_epoch = 0  # 从第0个epoch开始

    # 继续训练剩余的epochs
    for epoch in range(start_epoch, total_epochs):
        print(f"开始训练epoch {epoch}...")
        time.sleep(1)  # 模拟训练过程，暂停1秒
        print(f"完成训练epoch {epoch}.")

        # 每个epoch结束后保存权重和epoch信息
        save_weights(weights, epoch, weights_file)  # 保存当前的权重和epoch信息

# 如果该脚本作为主程序运行，则执行main函数
if __name__ == '__main__':
    main()