TensorFlow2.0:模型的保存与加载

梅文化_2019

于 2019-09-11 09:33:28 发布

阅读量2.8k

点赞数 3

CC 4.0 BY-SA版权

分类专栏： tensorflow2.0

本文链接：https://blog.youkuaiyun.com/meijie2018_1/article/details/100726778

本文详细介绍了在TensorFlow2.0中如何进行权重参数和完整模型的保存与加载。针对权重参数的加载，强调了模型结构需要与原始模型一致。而对于完整模型的保存与加载，说明了此方法适用于Sequential模型，不适用于自定义模型。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一、权重参数的保存与加载

network.save_weights('weights.ckpt')
network.load_weights('weights.ckpt')

权重参数的保存与加载可以针对任何模型，包括自定义的。
但是在加载权重参数时，其模型的结构需要与原来的完全一致。

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers,Sequential,optimizers,datasets,metrics

def preprocess(x,y):
    x = tf.cast(tf.reshape(x,[-1]),dtype=tf.float32)/255.
    y = tf.cast(tf.one_hot(y,depth=10),dtype=tf.int32)
    return x,y

#load_data
(x_train,y_train),(x_val,y_val) = datasets.mnist.load_data()
print('data: ',x_train.shape,y_train.shape,x_val.shape,y_val.shape)

db = tf.data.Dataset.from_tensor_slices((x_train,y_train))
db = db.map(preprocess).shuffle(60000).batch(128)
db_val = tf.data.Dataset.from_tensor_slices((x_val,y_val))
db_val = db_val.map(preprocess).batch(128)

#self def layer
class MyDense(layers.Layer):#inhe

最低0.47元/天解锁文章