Изменить модель в тензор-федерированном, но не работать

Question

Я пытаюсь изменить модель (просто и скрытый слой) в учебнике по Федеративному обучению для классификации изображений. Но результат показывает, что w1 и b1 не меняются и сохраняют начальное значение 0 после нескольких итераций. Только w2 и b2 обучаемы на тренировке. Вот мой код:

MnistVariables = collections.namedtuple(
    'MnistVariables', 'w1 w2 b1 b2 num_examples loss_sum accuracy_sum')


def create_mnist_variables():
    return MnistVariables(
        w1=tf.Variable(
            lambda: tf.zeros(dtype=tf.float32, shape=(784, 128)),
            name='w1',
            trainable=True),
        w2=tf.Variable(
            lambda: tf.zeros(dtype=tf.float32, shape=(128, 10)),
            name='w2',
            trainable=True),
        b1=tf.Variable(
            lambda: tf.zeros(dtype=tf.float32, shape=(128)),
            name='b1',
            trainable=True),
        b2=tf.Variable(
            lambda: tf.zeros(dtype=tf.float32, shape=(10)),
            name='b2',
            trainable=True),
        num_examples=tf.Variable(0.0, name='num_examples', trainable=False),
        loss_sum=tf.Variable(0.0, name='loss_sum', trainable=False),
        accuracy_sum=tf.Variable(0.0, name='accuracy_sum', trainable=False))


def mnist_forward_pass(variables, batch):
    a = tf.add(tf.matmul(batch['x'], variables.w1) , variables.b1)
    a= tf.nn.relu(a)
    y = tf.nn.softmax(tf.add(tf.matmul(a, variables.w2),variables.b2))
    predictions = tf.cast(tf.argmax(y, 1), tf.int32)

    flat_labels = tf.reshape(batch['y'], [-1])
    loss = -tf.reduce_mean(tf.reduce_sum(
        tf.one_hot(flat_labels, 10) * tf.log(y), reduction_indices=[1]))
    accuracy = tf.reduce_mean(
        tf.cast(tf.equal(predictions, flat_labels), tf.float32))

    num_examples = tf.to_float(tf.size(batch['y']))

    tf.assign_add(variables.num_examples, num_examples)
    tf.assign_add(variables.loss_sum, loss * num_examples)
    tf.assign_add(variables.accuracy_sum, accuracy * num_examples)

    return loss, predictions


def get_local_mnist_metrics(variables):
    return collections.OrderedDict([
        ('w1', variables.w1),
        ('w2', variables.w2),
        ('b1', variables.b1),
        ('b2', variables.b2),
        ('num_examples', variables.num_examples),
        ('loss', variables.loss_sum / variables.num_examples),
        ('accuracy', variables.accuracy_sum / variables.num_examples)
        ])

class MnistModel(tff.learning.Model):

    def __init__(self):
        self._variables = create_mnist_variables()

    @property
    def trainable_variables(self):
        return [self._variables.w1, self._variables.w2,
                self._variables.b1, self._variables.b2]

Я также добавляю w2 и b2 в обучаемые переменные. Но кажется, что они не обучены в процессе обучения, и я не знаю почему. У кого-нибудь есть успешный опыт изменения модели в этом уроке?

Answer 1

Я подозреваю, что активация ReLU с нулевой инициализацией w1 а также b1 проблематичны, и это может быть случай "умирающего ReLU" (см. Что такое проблема "умирающего ReLU" в нейронных сетях?

поскольку w1 а также b1 инициализируются на ноль, я ожидаю, что выход также будет 0 после матрицы умножение и сложение.

Возможные варианты: использовать ненулевой инициализатор, использовать альтернативную функцию активации (или не иметь активации после первого уровня).