0. 概述

在pytorch中有两种方式可以保存推理模型，第一种是只保存模型的参数，比如parameters和buffers；另外一种是保存整个模型；

1. 保存模型 - 权重参数

我们可以用torch.save()函数来保存model.state_dict();state_dict()里面包含模型的parameters&buffers;这种方法只保存模型中必要的训练参数。你可以用pytorch中的pickle来保存模型；使用这种方法可以生成最直观的语法，并涉及最少的代码；这种方法的缺点是，序列化的数据被绑定到特定的类和保存模型时使用的确切的目录结构。这样做的原因是pickle并不保存模型类本身。相反，它保存包含类的文件的路径，在加载期间使用；因此，当在其他项目中使用或重构后，您的代码可能以各种方式中断。我们将探讨如何保存和加载模型进行推断的两种方法。

• 步骤：
  （1）导入所有必要的库来加载我们的数据
  （2）定义和初始化神经网络
  （3）初始化优化器
  （4）保存并通过state_dict加载模型
  （5）保存并加载整个模型

1.1 代码

# -*- coding: utf-8 -*-
# @Project: zc
# @Author: zc
# @File name: Neural_Network_test
# @Create time: 2022/3/19 15:33

# 1.导入相关数据库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn import functional as F


# 2.定义神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# 3. 实例化神经网络
net = Net()

# 4. 实例化优化器
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 5. 保存模型参数
# Specify a path
PATH = "state_dict_model.pt"

# 6. 保存模型的参数字典：parameters and buffers
torch.save(net.state_dict(), PATH)

# 7. 实例化新的模型
model = Net()

# 8. 给新的实例加载之前的模型参数
model.load_state_dict(torch.load(PATH))

# 9. 设置模型为评估模式
model.eval()

• 注意（1）：
  pytorch中常用的惯例是将model.state_dict()保存为"state_dict_model.pt"，即文件的格式一般是.pt或者.pth格式文件；注意load_state_dict加载的是一个字典，而不是路径。
• 注意（2）：
  模型参数在推理阶段一定要设置model.eval();这样可以让dropout和batchnorm失效，如果没设置推理模式，会得到不一样的结果。

2. 保存模型 - 整个模型

将模型所有的内容都保存下来。

# Specify a path
PATH = "entire_model.pt"

# Save
torch.save(net, PATH)

# Load
model = torch.load(PATH)
model.eval()

3. 保存模型 - checkpoints

我们按照checkpoints模式来保存模型，本质上就是按照字典的模式进行分门别类的保存，我们可以通过键值进行加载。

• epoch：训练周期
• model_state_dict：模型可训练参数
• optimizer_state_dict：模型优化器参数
• loss:模型的损失函数

# Additional information
EPOCH = 5
PATH = "model.pt"
LOSS = 0.4

torch.save({
            'epoch': EPOCH,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS,
            }, PATH)

保存和加载通用的检查点模型以进行推断或恢复训练，这有助于您从上一个地方继续进行。当保存一个常规检查点时，您必须保存模型的state_dict之外的更多信息。保存优化器的state_dict也很重要，因为它包含缓冲区和参数，随着模型的运行而更新。您可能希望保存的其他项目是您离开的时期，最新记录的训练损失，外部torch.nn.嵌入层，以及更多，基于自己的算法

3.1 代码

# 1.导入相关数据库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn import functional as F

# 2. 定义神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 3. 实例化神经网络
net = Net()

# 4. 实例化优化器
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# Additional information

# 5. 定义超参数
EPOCH = 5
PATH = "model.pt"
LOSS = 0.4

# 6. 以checkpoints形式保存模型的相关数据
torch.save({
            'epoch': EPOCH,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS,
            }, PATH)

# 7. 重新实例化一个模型
model = Net()

# 8. 实例化优化器
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)


# 9. 加载以前的checkpoint
checkpoint = torch.load(PATH)

# 10. 通过键值来加载相关参数
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

# 11.设置推理模式
model.eval()
# - or -
model.train()

4. 保存双模型

当保存有多个神经网络模型组成的神经网络时，比如GAN对抗模型，sequence-to-sequence序列到序列模型，或者一个组合模型，你必须为每一个模型保存状态字典state_dict()和其对应的优化器参数optimizer.state_dict();您还可以保存任何其他项目，可能会帮助您恢复训练，只需将它们添加到字典；为了加载模型，第一步是初始化神经网络模型和优化器，然后用torch.load()去加载checkpoint对应的数据,因为checkpoints是字典，所以我们可以通过键值进行查询导入；

4.1 相关步骤

（1）导入所有相关的数据库
（2）定义和实例化神经网络模型
（3）初始化优化器
（4）保存多重模型
（5）加载多重模型

# 1.导入相关数据库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn import functional as F

# 2. 定义神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 3. 实例化神经网络A,B
netA = Net()
netB = Net()

# 4. 实例化优化器A,B
optimizerA = optim.SGD(netA.parameters(), lr=0.001, momentum=0.9)
optimizerB = optim.SGD(netB.parameters(), lr=0.001, momentum=0.9)

# 5. 保存模型
# Specify a path to save to
PATH = "model.pt"

torch.save({
            'modelA_state_dict': netA.state_dict(),
            'modelB_state_dict': netB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            }, PATH)

# 6.重新实例化新的网络模型A,B
modelA = Net()
modelB = Net()

# 7. 重新实例化新的网络模型A,B
optimModelA = optim.SGD(modelA.parameters(), lr=0.001, momentum=0.9)
optimModelB = optim.SGD(modelB.parameters(), lr=0.001, momentum=0.9)

# 8. 将以前模型的参数重新加载到新的模型A,B中
checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

# 9. 开启预测模式
modelA.eval()
modelB.eval()
# - or -
# 10.开启训练模式
modelA.train()
modelB.train()

5. 机器学习流程图

6. 机器学习常用库

Keras训练模型有多种保存方法，可以保存为hdf5文件，也可以保存为json格式文件，可以同时保存模型图和权重，也可以单独保存模型图和权重，还可以保存为tensorflow-serving支持的pb格式。下面以一个简单的模型分别来介绍不同的保存方法。

模型图构建

下面用keras中函数式API构建一个简单的LSTM多分类模型，模型具体结构如下：

import keras
from keras.models import Sequential
from keras.layers import Input, Dense, Dropout, Embedding, LSTM
from keras.models import Model
from keras import backend as K
K.clear_session()

maxlen = 50
vocab_size = 1024
embedding_size = 128
hidden_size = 128
num_classes = 10

inputs = Input(shape=(maxlen,), name="inputs")
x = Embedding(input_dim=vocab_size, output_dim=embedding_size)(inputs)
x = LSTM(units=hidden_size)(x)
x = Dropout(rate=0.5)(x)
outputs = Dense(num_classes, activation='softmax', name="outputs")(x)
model = Model(inputs=inputs, outputs=outputs)

print(model.summary())

模型结构图如下：

训练模型

这里只是一个训练示例，为了减小数据预处理的过程，我们直接用随机产生的虚拟数据。

# 生成虚拟数据
import numpy as np

data = np.random.randint(vocab_size, size=(1000, maxlen))
labels = np.random.randint(num_classes, size=(1000, 1))
onehot_labels = keras.utils.to_categorical(labels, num_classes=10)

x_train, y_train = data[:800], onehot_labels[:800]
x_test, y_test = data[800:], onehot_labels[800:]

model.compile(loss='categorical_crossentropy',
              optimizer='rmsprop',
              metrics=['accuracy'])
model.fit(x_train, y_train, batch_size=32, epochs=10)
score, acc = model.evaluate(x_test, y_test, batch_size=32)
print('Test score: %s, Test acc: %s' % (score, acc))

保存模型

下面分别介绍保存为hdf5格式，json格式，以及tf-serving中的pb格式。

hdf5格式持久化

keras模型保存为hdf5又分为同时保存模型图和权重，只保存权重不保存模型图

一、同时保存模型图和权重

import keras
from keras.models import load_model

# 保存训练好的model为hdf5文件
model.save('model/my_model.h5')  
# 重新加载模型
model = load_model('model/my_model.h5')

二、只保存权重不保存模型图

# 只保存模型权重
model.save_weights('model/my_model_weights.h5')

# 加载模型权重，这里加载权重前需要重新构建model的模型图
model = Model(inputs=..., outputs=...)
model.load_weights('model/my_model_weights.h5')

json格式持久化

我们尝试把模型保存为json格式，模型图和权重可以单独保存，也可以同时保存。

一、json格式模型图

import keras
from keras.models import model_from_json

# 模型图保存为json格式字符串
model_json = model.to_json()
# 从json字符串加载模型
model = model_from_json(model_json)

如果需要持久化保存加入json的保存和读取操作

import keras
from keras.models import model_from_json

model_json = model.to_json()

# json格式模型图保存到文件
json_dict = json.loads(model_json)
with open("model/model_graph.json", "w") as fw:
    json.dump(json_dict, fw, indent=4)
    
# 加载json格式模型图
with open("model/model_graph.json", "r") as fr:
    param_data = json.load(fr)
model_json = json.dumps(param_data)

model = model_from_json(model_json)

二、json格式保存权重
keras中并没有直接把权重保存成json格式的方法，但是有获取参数名和权重的方法，调用model.get_weights()返回模型中所有权重张量的列表，类型为 Numpy 数组。

1. 简单的方法，直接调用get_weights和set_weights方法。

# 获取模型权重
weights = model.get_weights() 
# 给模型权重赋值，注意列表中的数组必须与 get_weights() 返回的权重具有相同的尺寸。
model.set_weights(weights)

2. 持久化的方法，没有详细测试，可能有些模型不行。model.layers可以获得模型的所有层，layer.weights可以获得层内参数的信息，layer.get_weights()可以获得层内每个参数的权重值。

import json

# 获得模型参数和对应权重
weights_dict = {}
for layer in model.layers:
    for weight, value in zip(layer.weights, layer.get_weights()):
        weight_value = value.tolist()
        weights_dict[weight.name] = weight_value

# 模型权重保存为json格式
with open("model/model_weights.json", "w") as fw:
    json.dump(weights_dict, fw, indent=4)

tf-serving调用的pb格式

keras模型也可以转成tf-serving调用的pb格式，利用tf-serving加载模型可以同时支持grpc和Restful API调用。模型转化方法如下：

import os
import tensorflow as tf
from keras import backend as K
from keras.models import load_model
from keras.models import Model

def export_model(model, export_model_dir, model_version):
    with tf.get_default_graph().as_default():
        # prediction_signature
        tensor_info_input = tf.saved_model.utils.build_tensor_info(model.input)
        tensor_info_output = tf.saved_model.utils.build_tensor_info(model.output)
        
        prediction_signature = (
            tf.saved_model.signature_def_utils.build_signature_def(
                inputs={'inputs': tensor_info_input}, # Tensorflow.TensorInfo
                outputs={'outpus': tensor_info_output},
                method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
        )

        # set-up a builder
        os.mkdir(export_model_dir)
        export_path = os.path.join(tf.compat.as_bytes(export_model_dir), tf.compat.as_bytes(str(model_version)))
        builder = tf.saved_model.builder.SavedModelBuilder(export_path)
        builder.add_meta_graph_and_variables(
            sess=K.get_session(),
            tags=[tf.saved_model.tag_constants.SERVING],
            signature_def_map={
                'predict':prediction_signature,
                tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: prediction_signature,
            },
            )
        builder.save()

model = load_model("model/model_save.h5")
export_model(model, "test_model", "1")

参看文档：
[1]. https://keras.io/zh/models/about-keras-models/
[2]. https://keras.io/zh/getting-started/faq/
[3]. https://www.tensorflow.org/tfx/tutorials/serving/rest_simple
[4]. https://www.tensorflow.org/tfx/serving/serving_basic

一、保存模型的全部配置信息

使用model.save()函数搭配tf.keras.models.load_model()对模型的架构，权重以及配置进行保存与恢复。

模型的保存代码如下：

import tensorflow as tf
import os

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['acc']
)

# 模型的训练
model.fit(train_images, train_labels, epochs=3)

# 模型的评测
print(model.evaluate(test_images, test_labels, verbose=0))

# 模型的保存
model.save(r'model_data/model.h5')

模型的恢复代码如下：

import tensorflow as tf
import os

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的恢复
model = tf.keras.models.load_model(r'model_data/model.h5')

# 模型的评测
print(model.evaluate(test_images, test_labels, verbose=0))

二、仅仅保存模型的架构

使用model.to_json()搭配tf.keras.models.model_from_json()对模型的架构进行保存与恢复

模型的保存代码如下：

import tensorflow as tf
import os
import json

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['acc']
)

# 模型的训练
model.fit(train_images, train_labels, epochs=3)

# 模型的评测
print(model.evaluate(test_images, test_labels, verbose=0))

# 模型的保存
# 生成json文件
model_json = model.to_json()
# 写入json文件
with open(r'model_json.json', 'w') as f:
    json.dump(model_json, f)
    print('模型的架构json文件保存完成！')

模型的恢复代码如下：

import tensorflow as tf
import os
import json

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的恢复
# 读取json文件
with open(r'model_json.json', 'r') as f:
    model_json = json.load(f)
# 模型的加载
model = tf.keras.models.model_from_json(model_json)
model.summary()

三、仅仅保存模型的权重

使用model.save_weights()函数搭配model.load_weights()函数对模型进行权重的保存与加载

模型权重的保存代码如下：

import tensorflow as tf
import os

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['acc']
)

# 模型的训练
model.fit(train_images, train_labels, epochs=3)

# 模型的评测
print(model.evaluate(test_images, test_labels, verbose=0))

# 模型的保存
model.save_weights(r'model_data/save_weights.h5')

模型权重的恢复代码如下：

import tensorflow as tf
import os

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['acc']
)

# 模型加载前的的评测
print('权重加载前的模型准确率')
print(model.evaluate(test_images, test_labels, verbose=0))
# 模型权重的加载
model.load_weights(r'model_data/save_weights.h5')
# 模型加载后的的评测
print('权重加载后的模型准确率')
print(model.evaluate(test_images, test_labels, verbose=0))

模型权重加载前后对比效果图

四、使用回调函数对模型进行保存

使用tf.keras.callbacks.ModelCheckpoint()回调函数对模型进行保存

模型权重的保存代码如下：

import tensorflow as tf
import os

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['acc']
)

# 设置回调函数
cp_callback = tf.keras.callbacks.ModelCheckpoint(r'checkpoint_data/logs',
                                                 save_weights_only=True)

# 模型的训练
model.fit(train_images, train_labels, epochs=3, callbacks=[cp_callback])

# 模型的评测
print(model.evaluate(test_images, test_labels, verbose=0))

模型的恢复代码如下：

import tensorflow as tf
import os

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['acc']
)

# 模型加载前的准确率
print(model.evaluate(test_images, test_labels, verbose=0))

# 模型的恢复
model.load_weights(r'checkpoint_data/logs')

# 模型加载后的准确率
print(model.evaluate(test_images, test_labels, verbose=0))

模型权重的恢复前后的对比图

五、对于自定义训练的模型进行保存

使用tf.train.Checkpoint()函数的.save()方法与.restore()方法对模型进行保存与恢复

模型的保存代码如下：

import tensorflow as tf
import os
import tqdm

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()
dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
dataset = dataset.batch(60000)
test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(60000)

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
optimizer = tf.keras.optimizers.Adam()
loss_func = tf.keras.losses.SparseCategoricalCrossentropy()

train_loss_mean = tf.keras.metrics.Mean('train_loss')
train_accuracy = tf.keras.metrics.Accuracy('train_accuracy')

test_loss_mean = tf.keras.metrics.Mean('test_loss')
test_accuracy = tf.keras.metrics.Accuracy('test_accuracy')

# 定义模型保存的函数
checkpoint = tf.train.Checkpoint(model=model)


# 定义单步的训练
def step_train(mol, images, labels):
    with tf.GradientTape() as t:
        pre = mol(images)
        loss_step = loss_func(labels, pre)
    grads = t.gradient(loss_step, mol.trainable_variables)
    optimizer.apply_gradients(zip(grads, mol.trainable_variables))
    train_loss_mean(loss_step)
    train_accuracy(labels, tf.argmax(pre, axis=-1))


def step_test(mol, imags, labels):
    pre = mol(imags, training=False)
    loss_step = loss_func(labels, pre)
    test_loss_mean(loss_step)
    test_accuracy(labels, tf.argmax(pre, axis=-1))


# 定义训练函数
def train():
    for i in range(300):
        tqdm_train = tqdm.tqdm(iter(dataset), total=len(dataset))
        for img, lab in tqdm_train:
            step_train(model, img, lab)
            tqdm_train.set_description_str('Epoch : {:3}'.format(i))
            tqdm_train.set_postfix_str(
                'train loss is {:.14f} train accuracy is {:.14f}'.format(train_loss_mean.result(),
                                                                         train_accuracy.result()))
        tqdm_test = tqdm.tqdm(iter(test_dataset), total=len(test_dataset))
        for ima, lbl in tqdm_test:
            step_test(model, ima, lbl)
            tqdm_test.set_description_str('Epoch : {:3}'.format(i))
            tqdm_test.set_postfix_str(
                'test loss is {:.14f} test accuracy is {:.14f}'.format(test_loss_mean.result(), test_accuracy.result()))
        if i % 50 == 0:
            checkpoint.save(file_prefix=r'save_check/logs')
        train_loss_mean.reset_states()
        train_accuracy.reset_states()
        test_loss_mean.reset_states()
        test_accuracy.reset_states()
        tqdm_train.close()
        tqdm_test.close()


if __name__ == '__main__':
    train()

模型的恢复代码如下：

import tensorflow as tf
import os
import tqdm

# 环境变量的配置
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

# 数据的加载
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.fashion_mnist.load_data()
test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(60000)

# 模型的构建
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 模型的相关配置
test_accuracy = tf.keras.metrics.Accuracy('test_accuracy')

# 定义模型保存的函数
checkpoint = tf.train.Checkpoint(model=model)
checkpoint.restore(tf.train.latest_checkpoint(r'save_check'))


def step_test(mol, imags, labels):
    pre = mol(imags, training=False)
    test_accuracy(labels, tf.argmax(pre, axis=-1))


tqdm_test = tqdm.tqdm(iter(test_dataset), total=len(test_dataset))

for img, lable in tqdm_test:
    step_test(model, img, lable)
    tqdm_test.set_postfix_str(test_accuracy.result())