ester
ester

Reputation: 11

change model in tensorflow-federated but not work

I try to change model(just and hidden layer) in the tutorial of Federated Learning for Image Classification. But the result shows that w1 and b1 don't change and retain the initial value 0 after multiple iterations. Only w2 and b2 are trainable in the training. Here is my code:

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]

I also add w2 and b2 in the trainable variables. But it seems that they are not trained in the training process and I don't know why. Does anyone have some successful experiences to change model in this tutorial?

Upvotes: 0

Views: 111

Answers (1)

Zachary Garrett
Zachary Garrett

Reputation: 2941

I suspect the ReLU activations with zero initialisations of w1 and b1 are problematic and this maybe a case of "dying ReLU" (see What is the “dying ReLU” problem in neural networks?.

Since w1 and b1 are initialized to zero, I would expect the output to also be 0 after the matrix multiply and addition.

Possible options to try: using a non-zero initializer, use an alternative activation function (or don't have an activation after the first layer).

Upvotes: 1

Related Questions