TensorFlow优化器梯度步长缩放：基于Master网络权重训练Worker网络

这是个很巧妙的加权训练策略——让Master网络动态调整Worker的梯度更新强度，本质上是让表现更好的Worker获得更大的训练权重。我会结合tf.layers和Adam优化器，一步步给你实现这个逻辑：

1. 先定义Worker和Master的网络结构 #

首先用tf.layers搭建你的Worker网络（这里做了简化，你可以根据实际任务修改），然后搭建Master网络来输出每个Worker的权重：

import tensorflow as tf

def build_worker_network(inputs, num_classes):
    """构建单个Worker网络"""
    x = tf.layers.dense(inputs, 256, activation=tf.nn.relu)
    x = tf.layers.dense(x, 128, activation=tf.nn.relu)
    logits = tf.layers.dense(x, num_classes)
    return logits

def build_master_network(worker_logits_list, num_workers):
    """构建Master网络，输出每个Worker的权重（归一化到[0,1]）"""
    # 把所有Worker的logits拼接作为Master的输入（你也可以用Worker的其他特征，比如准确率）
    concatenated_features = tf.concat(worker_logits_list, axis=1)
    x = tf.layers.dense(concatenated_features, 64, activation=tf.nn.relu)
    # 用softmax确保权重和为1，符合投票权重的逻辑
    worker_weights = tf.layers.dense(x, num_workers, activation=tf.nn.softmax)
    return worker_weights

2. 构建训练图：损失计算+梯度缩放核心逻辑 #

接下来是关键部分：计算每个Worker的损失，用Master输出的权重缩放梯度，再交给Adam优化器更新参数。另外，Master网络也需要自己的损失来学习合理的权重，这里我们让Master的权重和Worker的预测准确率对齐：

# 超参数设置
num_workers = 3
num_classes = 10
input_dim = 784  # 示例用MNIST输入维度
learning_rate = 1e-3

# 输入占位符
inputs = tf.placeholder(tf.float32, shape=[None, input_dim], name="inputs")
labels = tf.placeholder(tf.int32, shape=[None], name="labels")

# 构建所有Worker网络，计算各自的损失
worker_logits = []
worker_losses = []
worker_predictions = []
for _ in range(num_workers):
    logits = build_worker_network(inputs, num_classes)
    worker_logits.append(logits)
    # Worker的交叉熵损失
    loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=labels, logits=logits
    ))
    worker_losses.append(loss)
    # 记录Worker的预测结果，用于训练Master
    pred = tf.argmax(logits, axis=1)
    worker_predictions.append(pred)

# 构建Master网络，得到每个Worker的权重
worker_weights = build_master_network(worker_logits, num_workers)

# ---------------------- 核心：用Master权重缩放Worker的梯度 ----------------------
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)

# 1. 计算每个Worker的原始梯度
worker_grads_list = []
for loss in worker_losses:
    # compute_gradients会返回(梯度, 变量)的列表
    grads_and_vars = optimizer.compute_gradients(loss)
    worker_grads_list.append(grads_and_vars)

# 2. 用Master的权重缩放每个Worker的梯度
scaled_grads_list = []
for worker_idx in range(num_workers):
    # 取当前Worker对应的权重，这里对batch内的权重取均值作为全局缩放因子
    # 如果需要按样本缩放，可以直接用worker_weights[:, worker_idx]与梯度相乘（注意维度匹配）
    weight = tf.reduce_mean(worker_weights[:, worker_idx])
    scaled_grads = []
    for grad, var in worker_grads_list[worker_idx]:
        if grad is not None:  # 跳过没有梯度的变量（比如共享参数可能重复？）
            scaled_grad = grad * weight
            scaled_grads.append((scaled_grad, var))
        else:
            scaled_grads.append((grad, var))
    scaled_grads_list.append(scaled_grads)

# ---------------------- 训练Master网络：让权重匹配Worker的准确率 ----------------------
# 计算每个Worker的预测正确率
worker_correct = [tf.cast(tf.equal(pred, labels), tf.float32) for pred in worker_predictions]
# Master的损失：让输出的权重和Worker的实际正确率尽可能接近（MSE损失）
master_loss = tf.reduce_mean(tf.square(worker_weights - tf.stack(worker_correct, axis=1)))
# 计算Master网络的梯度
master_grads = optimizer.compute_gradients(master_loss)

# ---------------------- 合并所有梯度，生成训练操作 ----------------------
# 把所有Worker的缩放梯度和Master的梯度合并
all_grads = []
for grads in scaled_grads_list:
    all_grads.extend(grads)
all_grads.extend(master_grads)

# 应用所有梯度更新参数
train_op = optimizer.apply_gradients(all_grads)

3. 关键细节说明 #

• 梯度缩放的灵活调整：上面示例用了batch权重的均值做全局缩放，如果你的Master输出是每个样本对应每个Worker的权重，你可以直接在损失计算时就给每个样本的损失乘以对应权重（loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(...) * worker_weights[:, worker_idx])），这样梯度会自动带上样本级的权重，可能更贴合你的需求。
• 共享参数处理：如果多个Worker共享部分参数，要注意梯度的累加——上面的代码会自动把不同Worker对同一参数的缩放梯度相加，这是合理的，相当于共享参数的更新是所有Worker加权后的梯度总和。
• 权重约束：用softmax输出Master的权重确保了权重和为1，你也可以用sigmoid让每个权重独立在[0,1]区间，根据你的投票逻辑选择即可。

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