TensorFlow 1.x Keras模型保存报错及Web部署方案咨询

参考公开代码实现Deep Ranking模型训练过程中，模型保存环节执行失败，初步判断Lambda层为导致保存失败的原因，结合技术社区相关讨论可知，TensorFlow生态下部分不可序列化对象会触发模型保存报错。

相关实现代码如下：

from __future__ import print_function

import numpy as np
from keras.applications.vgg16 import VGG16
from keras.layers import *
from keras.models import Model, load_model
from keras.optimizers import SGD
# from keras.optimizers import gradient_descent_v2 
from keras.preprocessing.image import load_img, img_to_array
import tensorflow.compat.v1 as tf
from keras import backend as K


config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
K.set_session(sess)

def convnet_model_():
    vgg_model = VGG16(weights=None, include_top=False)
    x = vgg_model.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(4096, activation='relu')(x)
    x = Dropout(0.6)(x)
    x = Dense(4096, activation='relu')(x)
    x = Dropout(0.6)(x)
    # 注意：此处Lambda层存在变量引用错误
    x = Lambda(lambda x_: K.l2_normalize(x, axis=1))(x)
    convnet_model = Model(inputs=vgg_model.input, outputs=x)
    return convnet_model


def deep_rank_model():
    convnet_model = convnet_model_()
    first_input = Input(shape=(224, 224, 3))
    first_conv = Conv2D(96, kernel_size=(8, 8), strides=(16, 16), padding='same')(first_input)
    first_max = MaxPool2D(pool_size=(3, 3), strides=(4, 4), padding='same')(first_conv)
    first_max = Flatten()(first_max)
    first_max = Lambda(lambda x: K.l2_normalize(x, axis=1))(first_max)

    second_input = Input(shape=(224, 224, 3))
    second_conv = Conv2D(96, kernel_size=(8, 8), strides=(32, 32), padding='same')(second_input)
    second_max = MaxPool2D(pool_size=(7, 7), strides=(2, 2), padding='same')(second_conv)
    second_max = Flatten()(second_max)
    second_max = Lambda(lambda x: K.l2_normalize(x, axis=1))(second_max)

    merge_one = concatenate([first_max, second_max])

    merge_two = concatenate([merge_one, convnet_model.output])
    emb = Dense(4096)(merge_two)
    l2_norm_final = Lambda(lambda x: K.l2_normalize(x, axis=1))(emb)

    final_model = Model(inputs=[first_input, second_input, convnet_model.input], outputs=l2_norm_final)

    return final_model

执行模型保存操作时，调用save_weights保存权重成功，但调用to_json()导出模型结构环节抛出如下异常：

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-20-bdb0097d23ba> in <module>()
      3 deep_rankmodel.save_weights(model_path)
      4 f = open('deepranking.json','w')
----> 5 f.write(deep_rankmodel.to_json())
      6 f.close()

29 frames
/usr/lib/python3.7/copy.py in deepcopy(x, memo, _nil)
    167                     reductor = getattr(x, "__reduce_ex__", None)
    168                     if reductor:
---> 169                         rv = reductor(4)
    170                     else:
    171                         reductor = getattr(x, "__reduce__", None)

TypeError: can't pickle _thread.RLock objects

调试界面截图：

咨询问题共2个：

1. 上述报错的可行解决方案是什么
2. 是否存在无需依赖导出的JSON模型文件即可在服务端加载模型的方案，是否可以通过JavaScript等技术栈在服务端重构模型实现Web端部署

首先先修正代码里的显性逻辑错误：convnet_model_函数中第一个L2归一化Lambda层的入参写反了，lambda定义的形参是x_，内部却引用了外层作用域的x张量，哪怕不触发保存错误，推理结果也完全不对，需要先把这行改成x = Lambda(lambda x_: K.l2_normalize(x_, axis=1))(x)。

触发can't pickle _thread.RLock objects的核心原因是：代码使用tf.compat.v1静态图模式，手动创建的Session、后端张量属于带线程锁的不可序列化对象，被Lambda层的闭包上下文捕获后，to_json()执行深度拷贝序列化时无法pickle这类带锁对象，直接抛错。可落地的修复方案有3种：

• 方案1：替换Lambda层为官方内置可序列化层。删除所有封装L2归一化的Lambda层，直接用Keras内置的UnitNormalization(axis=1)实现L2归一化逻辑，内置层自带完整序列化配置，不会捕获外部不可序列化对象，改完后to_json()、权重保存均可正常执行。
• 方案2：移除手动创建Session的兼容写法。删除代码中手动创建tf.ConfigProto、tf.Session、K.set_session(sess)这三行代码，使用TensorFlow 2.x默认的Eager执行模式，避免静态图下的会话对象被Lambda层捕获。如果必须保留Lambda层，建议给每个Lambda层显式命名，减少匿名函数序列化时的上下文依赖。
• 方案3：更换模型保存方式。放弃to_json()+save_weights()的结构权重分离保存逻辑，直接调用model.save('deepranking.h5')保存整H5文件，或者用model.save('deepranking_savedmodel/')导出SavedModel格式，这两种保存方式会直接序列化计算图而非Python层对象，不会触发pickle线程锁的错误。

• 无JSON文件的服务端加载方案完全可行。直接使用上述提到的整H5或者SavedModel格式保存模型即可，加载时直接调用tf.keras.models.load_model(模型路径)就能同时还原模型结构和权重，完全不需要单独导出JSON结构文件。如果保留了自定义Lambda层，加载时只要在load_model中传入custom_objects参数声明Lambda层的逻辑即可正常加载。
• 基于JavaScript技术栈的Web部署也可实现，不需要用Python重构模型。流程为：先将训练完成的SavedModel模型通过TensorFlow JS提供的转换工具转成TFJS支持的格式（生成model.json结构文件+分片权重文件），之后无论是Node.js服务端还是浏览器端，都可以直接引入@tensorflow/tfjs库加载转换后的模型执行推理，不需要依赖Python运行环境。如果做服务端部署，Node.js加载TFJS模型的推理性能可以满足常规业务需求；如果做端侧部署，模型可以直接在用户浏览器中加载运行，不需要把图片等数据传回后端处理。

内容的提问来源于stack exchange，提问作者user14603474