【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()

代码解释

1. 导入必要的模块：

   import pickle
   import time
   import os
   import numpy as np
   

2. 定义权重计算函数：

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

3. 定义保存权重和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}.")
   

4. 定义从文件加载权重和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数据
   

5. 主程序逻辑：

   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()
   

该程序用于模拟长时间训练模型的过程，其中包括计算权重、训练模型、保存和加载权重以及当前训练的epoch。通过使用pickle模块，程序能够在中途停止并在重新启动时从上次停止的地方继续运行。