【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# 如果权重文件存在，则加载权重和epochif os.path.exists(weights_file):weights, start_epoch = load_weights(weights_file)  # 从文件加载权重和训练的epoch数else:# 否则，从第一个epoch开始，并计算权重weights = calculate_weights()  # 计算权重start_epoch = 0  # 从第0个epoch开始# 继续训练剩余的epochsfor 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)  # 如果存在，加载权重和epochelse:weights = calculate_weights()  # 如果不存在，计算权重start_epoch = 0  # 从第0个epoch开始# 继续训练剩余的epochsfor 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()