matwilso
matwilso

Reputation: 3214

How to use He initialization in TensorFlow

He / MSRA initialization, from Delving Deep into Rectifiers, seems to be a recommended weight initialization when using ReLUs.

Is there a built-in way to use this in TensorFlow? (similar to: How to do Xavier initialization on TensorFlow)?

Upvotes: 6

Views: 11483

Answers (1)

matwilso
matwilso

Reputation: 3214

TensorFlow 2.0

tf.keras.initializers.HeUniform()

or

tf.keras.initializers.HeNormal()

See docs for usage. (h/t to @mable)

TensorFlow 1.0

tf.contrib.layers.variance_scaling_initializer(dtype=tf.float32)

This will give you He / MRSA initialization. The documentation states that the default arguments for tf.contrib.layers.variance_scaling_initializer correspond to He initialization and that changing the arguments can yield Xavier initialization (this is what is done in TF's internal implementation for Xavier initialization).

Example usage:

W1 = tf.get_variable('W1', shape=[784, 256],
       initializer=tf.contrib.layers.variance_scaling_initializer())

or

initializer = tf.contrib.layers.variance_scaling_initializer()
W1 = tf.Variable(initializer([784,256]))

Upvotes: 11

Related Questions