如何在Xcode中修改.tflite模型内的条件实例归一化层参数？

我正在iOS应用中实现基于条件实例归一化（ConditionalInstanceNorm）的神经风格迁移模型，该模型支持单模型学习多风格并实现风格融合。Python端的Keras层实现如下：

class ConditionalInstanceNorm(tf.keras.layers.Layer):
  def __init__(self, scope_bn, y1, y2, alpha):
    super(ConditionalInstanceNorm, self).__init__()
    self.scope_bn = scope_bn
    self.y1 = y1
    self.y2 = y2
    self.alpha = alpha
  
  def build(self, input_shape):
    self.beta = self.add_weight(name="beta"+self.scope_bn, shape=(self.y1.shape[-1], input_shape[-1]), initializer=betaInitializer, trainable=True)
    self.gamma = self.add_weight(name="gamma"+self.scope_bn, shape=(self.y1.shape[-1], input_shape[-1]), initializer=gammaInitializer, trainable=True)
  
  def call(self, inputs):
    mean, var = tf.nn.moments(x=inputs, axes=[1,2], keepdims=True)
    beta1 = tf.matmul(self.y1, self.beta)
    gamma1 = tf.matmul(self.y1, self.gamma)
    beta2 = tf.matmul(self.y2, self.beta)
    gamma2 = tf.matmul(self.y2, self.gamma)
    beta = self.alpha*beta1 + (1. - self.alpha)*beta2
    gamma = self.alpha*gamma1 + (1. - self.alpha)*gamma2
    x = tf.nn.batch_normalization(x=inputs, mean=mean, variance=var, offset=beta, scale=gamma, variance_epsilon=1e-10)
    return x

在Python中，我可以通过遍历模型层修改y1、y2、alpha来切换风格组合：

for layer in filter(lambda x: "conditional_instance_norm" in x.name, model.layers):
    layer.y1 = y1
    layer.y2 = y2
    layer.alpha = alpha

但将模型转为TFLite后，在Swift中无法修改这些初始化时的参数，请问如何实现动态调整风格和融合权重？

TFLite是静态计算图，无法像Keras那样动态修改层的属性。要实现Swift端动态调整风格参数，需要将y1、y2、alpha改为模型的输入张量，而非层的初始化参数。具体步骤如下：

1. 修改Keras条件实例归一化层实现 #

重新定义层，把风格向量和融合权重作为call方法的输入，而非初始化参数：

class ConditionalInstanceNorm(tf.keras.layers.Layer):
  def __init__(self, scope_bn, style_dim):
    super(ConditionalInstanceNorm, self).__init__()
    self.scope_bn = scope_bn
    self.style_dim = style_dim  # 风格向量的维度
  
  def build(self, input_shape):
    # 根据风格向量维度定义beta和gamma的形状
    self.beta = self.add_weight(
        name="beta"+self.scope_bn,
        shape=(self.style_dim, input_shape[-1]),
        initializer="zeros",
        trainable=True
    )
    self.gamma = self.add_weight(
        name="gamma"+self.scope_bn,
        shape=(self.style_dim, input_shape[-1]),
        initializer="ones",
        trainable=True
    )
  
  def call(self, inputs, y1, y2, alpha):
    # inputs是特征图，y1/y2是风格向量，alpha是融合权重
    mean, var = tf.nn.moments(x=inputs, axes=[1,2], keepdims=True)
    beta1 = tf.matmul(y1, self.beta)
    gamma1 = tf.matmul(y1, self.gamma)
    beta2 = tf.matmul(y2, self.beta)
    gamma2 = tf.matmul(y2, self.gamma)
    
    # 广播alpha到匹配beta/gamma的形状
    alpha = tf.reshape(alpha, (-1, 1, 1, input_shape[-1]))
    beta = alpha * beta1 + (1. - alpha) * beta2
    gamma = alpha * gamma1 + (1. - alpha) * gamma2
    
    # 保持alpha的batch维度匹配输入
    beta = tf.reshape(beta, (-1, 1, 1, input_shape[-1]))
    gamma = tf.reshape(gamma, (-1, 1, 1, input_shape[-1]))
    
    x = tf.nn.batch_normalization(
        x=inputs,
        mean=mean,
        variance=var,
        offset=beta,
        scale=gamma,
        variance_epsilon=1e-10
    )
    return x

2. 重构模型结构 #

将风格向量y1、y2和融合权重alpha作为模型的额外输入，在调用ConditionalInstanceNorm层时传入：

def build_style_transfer_model(content_input_shape, style_dim=128):
    # 内容输入
    content_input = tf.keras.Input(shape=content_input_shape)
    # 风格输入：y1、y2是风格向量，alpha是融合权重
    y1_input = tf.keras.Input(shape=(style_dim,))
    y2_input = tf.keras.Input(shape=(style_dim,))
    alpha_input = tf.keras.Input(shape=(1,))
    
    # 示例编码器部分
    x = tf.keras.layers.Conv2D(64, (3,3), padding="same")(content_input)
    
    # 应用条件实例归一化层
    x = ConditionalInstanceNorm(scope_bn="bn1", style_dim=style_dim)(x, y1_input, y2_input, alpha_input)
    x = tf.keras.layers.Activation("relu")(x)
    
    # 后续层...（省略）
    
    # 输出转换后的图像
    output = tf.keras.layers.Conv2D(3, (3,3), padding="same", activation="tanh")(x)
    
    # 定义多输入模型
    model = tf.keras.Model(
        inputs=[content_input, y1_input, y2_input, alpha_input],
        outputs=output
    )
    return model

3. 转换为TFLite模型 #

转换时确保所有输入都被正确导出，使用TFLiteConverter：

model = build_style_transfer_model(content_input_shape=(256,256,3))
# 保存Keras模型
model.save("style_transfer_model.h5")

# 转换为TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 启用选择性量化（可选）
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# 保存TFLite模型
with open("style_transfer.tflite", "wb") as f:
    f.write(tflite_model)

4. Swift端推理实现 #

在iOS中，将内容图像、y1、y2、alpha作为输入张量传入模型：

import TensorFlowLite

// 加载TFLite模型
guard let modelPath = Bundle.main.path(forResource: "style_transfer", ofType: "tflite") else {
    fatalError("模型文件不存在")
}

var interpreter: Interpreter!
do {
    interpreter = try Interpreter(modelPath: modelPath)
    // 分配张量内存
    try interpreter.allocateTensors()
} catch {
    fatalError("初始化解释器失败: \(error)")
}

// 准备输入数据
// 1. 内容图像：转换为Float32张量，形状为[1, 256, 256, 3]
let contentImageTensor: Data = ... // 处理后的图像数据
// 2. 风格向量y1：形状为[1, 128]
let y1Data = Data(copyingBufferOf: styleVector1.map { Float32($0) })
// 3. 风格向量y2：形状为[1, 128]
let y2Data = Data(copyingBufferOf: styleVector2.map { Float32($0) })
// 4. 融合权重alpha：形状为[1, 1]
let alphaData = Data(copyingBufferOf: [Float32(0.5)])

// 设置输入张量
try interpreter.copy(contentImageTensor, toInputAt: 0)
try interpreter.copy(y1Data, toInputAt: 1)
try interpreter.copy(y2Data, toInputAt: 2)
try interpreter.copy(alphaData, toInputAt: 3)

// 运行推理
try interpreter.invoke()

// 获取输出张量
let outputTensor = try interpreter.output(at: 0)
let outputData = outputTensor.data
// 处理输出数据得到风格化图像...

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