<jupyter_start><jupyter_text>This notebook accompanies with my article [The effects of weight initialization on neural nets](https://www.wandb.com/articles/the-effects-of-weight-initialization-on-neural-nets).<jupyter_code># Selecting the right version of TensorFlow nad verifying it
%tensorflow_version 2.x
import tensorflow as tf
print(tf.__version__)
# Weights and Biases for experiment tracking
!pip install wandb
!wandb login
# Gather data and preprocess it
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()
X_train, X_test = X_train / 255.0, X_test / 255.0
CLASS_NAMES = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
import matplotlib.pyplot as plt

plt.figure(figsize=(10,10))
for i in range(25):
    plt.subplot(5,5,i+1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(True)
    plt.imshow(X_train[i], cmap=plt.cm.binary)
    plt.xlabel(CLASS_NAMES[y_train[i]])
plt.show()
import wandb
def get_training_model(init_scheme='zeros'):
    if isinstance(init_scheme, str):
        wandb.init(project='weight-initialization-tb', sync_tensorboard=True,
                   id=init_scheme)
    else:
        wandb.init(project='weight-initialization-tb', sync_tensorboard=True,
                   id=str(init_scheme))
    
    model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(256, activation='relu', kernel_initializer=init_scheme,
            bias_initializer='zeros'),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(128, activation='relu', kernel_initializer=init_scheme,
            bias_initializer='zeros'),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10, activation='softmax', kernel_initializer=init_scheme,
            bias_initializer='zeros')
    ])

    model.compile(optimizer='adam',
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])
    
    return model, wandb.run.dir
from wandb.keras import WandbCallback
model_w_zeros, run_dir = get_training_model()

tb_callback = tf.keras.callbacks.TensorBoard(log_dir=run_dir, histogram_freq=4, write_images=True)

model_w_zeros.fit(X_train, y_train, 
    validation_data=(X_test, y_test), 
    epochs=20, batch_size=128,
    callbacks=[WandbCallback(data_type='image', labels=CLASS_NAMES, 
                             validation_data=(X_test, y_test)),
               tb_callback])
model_w_ones, run_dir = get_training_model(init_scheme='ones')

tb_callback = tf.keras.callbacks.TensorBoard(log_dir=run_dir, histogram_freq=4, write_images=True)

model_w_ones.fit(X_train, y_train, 
    validation_data=(X_test, y_test), 
    epochs=20, batch_size=128,
    callbacks=[WandbCallback(data_type='image', labels=CLASS_NAMES, 
                             validation_data=(X_test, y_test)),
               tb_callback])
model_w_uniform, run_dir = get_training_model(init_scheme=tf.keras.initializers.RandomUniform(minval=0, maxval=1, seed=666))

tb_callback = tf.keras.callbacks.TensorBoard(log_dir=run_dir, histogram_freq=4, write_images=True)

model_w_uniform.fit(X_train, y_train, 
    validation_data=(X_test, y_test), 
    epochs=20, batch_size=128,
    callbacks=[WandbCallback(data_type='image', labels=CLASS_NAMES, 
                             validation_data=(X_test, y_test)),
               tb_callback])
model_w_centered, run_dir = get_training_model(init_scheme=tf.keras.initializers.RandomUniform(minval=-0.5, maxval=0.5, seed=666))

tb_callback = tf.keras.callbacks.TensorBoard(log_dir=run_dir, histogram_freq=4, write_images=True)

model_w_centered.fit(X_train, y_train, 
    validation_data=(X_test, y_test), 
    epochs=20, batch_size=128,
    callbacks=[WandbCallback(data_type='image', labels=CLASS_NAMES, 
                             validation_data=(X_test, y_test)),
               tb_callback])
def get_training_model_w_rule(init_scheme):
    if isinstance(init_scheme, str):
        wandb.init(project='weight-initialization-tb', sync_tensorboard=True,
                   id=init_scheme)
    else:
        wandb.init(project='weight-initialization-tb', sync_tensorboard=True,
                   id=str(init_scheme))
    
    model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    model.compile(optimizer='adam',
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])
    
    return model, wandb.run.dir
import numpy as np

model_w_uniform_rule, run_dir = get_training_model_w_rule('uniform-rule')

tb_callback = tf.keras.callbacks.TensorBoard(log_dir=run_dir, histogram_freq=4, write_images=True)

for layer in model_w_uniform_rule.layers:
    if isinstance(layer, tf.keras.layers.Dense):
        shape = (layer.weights[0].shape[0], layer.weights[0].shape[1])
        y = 1.0/np.sqrt(shape[0])
        rule_weights = np.random.uniform(-y, y, shape)
        layer.weights[0] = rule_weights
        layer.weights[1] = 0

model_w_uniform_rule.compile(optimizer='adam',
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])

model_w_uniform_rule.fit(X_train, y_train, 
    validation_data=(X_test, y_test), 
    epochs=20, batch_size=128,
    callbacks=[WandbCallback(data_type='image', labels=CLASS_NAMES, 
                             validation_data=(X_test, y_test)),
               tb_callback])
import numpy as np

model_w_normal_rule, run_dir = get_training_model_w_rule('normal-rule')

tb_callback = tf.keras.callbacks.TensorBoard(log_dir=run_dir, histogram_freq=4, write_images=True)

for layer in model_w_normal_rule.layers:
    if isinstance(layer, tf.keras.layers.Dense):
        shape = (layer.weights[0].shape[0], layer.weights[0].shape[1])
        y = 1.0/np.sqrt(shape[0])
        rule_weights = np.random.uniform(-y, y, shape)
        layer.weights[0] = rule_weights
        layer.weights[1] = 0

model_w_normal_rule.compile(optimizer='adam',
                loss='sparse_categorical_crossentropy',
                metrics=['accuracy'])

model_w_normal_rule.fit(X_train, y_train, 
    validation_data=(X_test, y_test), 
    epochs=20, batch_size=128,
    callbacks=[WandbCallback(data_type='image', labels=CLASS_NAMES, 
                             validation_data=(X_test, y_test)),
               tb_callback])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12387657279808269, 'fraction_numerical': 0.017675254643499103, 'mean_word_length': 3.5545702592087314, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '430839', 'license': 'Apache-2.0', 'n_tokens_mistral': 2364, 'n_tokens_neox': 2270, 'n_words': 324, 'path': 'Weights_Visualization_TensorBoard.ipynb', 'repo_name': 'joe733/TF-2.0-Hacks', 'size': '6676'}