tensorflow断点续训的方法

我用的 tensorflow 版本是 2.4.1

1.声明一个权重保存路径

model_wt_pth = r'./weights/model_wt.index'

---

2.模型初始化好之后先判断该路径是否存在，如果存在就加载该路径下的权重，然后再训练

model = Model()    # 初始化你的模型

if os.path.exists(model_wt_pth ):
    print('加载模型权重')
    model.load_weights(model_wt_pth )

---

3.训练之后将权重再保存到原路径中

if n % 500 == 0 and n != 0:    # n是训练次数，每500次训练后保存模型，训练次数自定义
    print('保存模型中...')
    model.save_weights(model_wt_pth .replace('.index', ''))    # 保存的时候不加后缀，2.4.1版的tensorflow
    print('保存完毕')