如何在TensorFlow Estimator中随输入传入模型参数？

刚好我之前在用TensorFlow Estimator的时候也碰到过这种自定义层参数传递的问题，给你分享几个实用的解决方案，你可以根据自己的场景选最合适的：

Estimator本身就提供了params参数来传递自定义配置，这是最规范的做法。你只需要在定义model_fn的时候加上params参数，然后在创建Estimator时把sigma放进params字典里就行。

示例代码如下：
首先修改你的model_fn，让它接收params：

def model_fun(features, labels, mode, params):
    # 从params字典里取出sigma
    sigma = params['sigma']
    # 接下来就可以用sigma来初始化你的自定义层了
    # ... 这里写你的模型构建逻辑 ...

然后创建Estimator的时候传入params：

classifier = tf.estimator.Estimator(
    model_fn=model_fun,
    params={'sigma': 0.3}  # 这里替换成你的sigma值
)

训练代码完全不用改，Estimator会自动把params传递给model_fn，非常省心。

如果你的sigma是固定不变的，不需要在训练过程中调整，那用functools.partial来包装model_fn会更简单。

首先导入functools模块，然后修改model_fn让它接收sigma参数：

import functools

def model_fun(features, labels, mode, sigma):
    # 直接使用传入的sigma构建层
    # ... 模型逻辑 ...

然后用partial把sigma的值固定下来，得到一个新的model_fn：

# 把sigma=0.5绑定到model_fun上
model_fn_with_sigma = functools.partial(model_fun, sigma=0.5)

# 创建Estimator时传入这个包装好的函数
classifier = tf.estimator.Estimator(model_fn=model_fn_with_sigma)

这种方法不需要改动Estimator的其他配置，适合sigma固定的场景。

如果你的sigma需要随每个训练批次变化（比如不同批次用不同的sigma值），那可以把sigma加到输入函数的特征字典里。

修改你的训练输入函数：

import numpy as np

train_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={
        "x": train_data,
        # 这里可以传入和train_data同长度的sigma数组，或者固定值广播
        "sigma": np.full(train_data.shape[0], 0.5)
    },
    batch_size=batch_size,
    num_epochs=nEpochs,
    shuffle=True
)

然后在model_fn里从features中取出sigma：

def model_fun(features, labels, mode):
    sigma = features['sigma']
    # 用这个sigma来构建你的自定义层
    # ... 模型逻辑 ...

这种方法适合需要动态调整sigma的场景，比如做参数敏感性实验的时候。

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