Visualizing the Loss Landscape of Neural Networks

Application showcasing landscapeviz here

This repository is an implementation of the paper

Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer and Tom Goldstein. Visualizing the Loss Landscape of Neural Nets. NIPS, 2018.

This code was implemented in tensorflow 2.0. The authors also have an implementation using pytorch.

How to use

# 1. define model

model = tf.keras.Sequential([
	tf.keras.layers.Dense(10, activation=tf.nn.relu, input_shape=(4,)),  # input shape required
	tf.keras.layers.Dense(10, activation=tf.nn.relu),
	tf.keras.layers.Dense(3, activation=tf.nn.softmax)
])

model.compile("sgd", loss="sparse_categorical_crossentropy", metrics=['sparse_categorical_accuracy', 'categorical_hinge'])

# 2. get data
data = sklearn.datasets.load_iris()
X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(data["data"], data["target"], test_size=0.25, random_state=seed)

scaler_x = sklearn.preprocessing.MinMaxScaler(feature_range=(-1,+1)).fit(X_train)
X_train = scaler_x.transform(X_train)
X_test = scaler_x.transform(X_test)

# 3. train model
model.fit(X_train, y_train, batch_size=32, epochs=60, verbose=0)


# 4. build mesh and plot
landscapeviz.build_mesh(model, (X_train, y_train), grid_length=40, verbose=0)
landscapeviz.plot_contour(key="sparse_categorical_crossentropy")
landscapeviz.plot_3d(key="sparse_categorical_crossentropy")