In [1]:
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
def train():
import numpy as np
import tensorflow as tf
import wandb
config_defaults = {
'hidden_nodes': 128
}
wandb.init(config=config_defaults)
fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
train_images = train_images / 255.0
test_images = test_images / 255.0
model = tf.keras.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(wandb.config.hidden_nodes, activation=tf.nn.relu),
tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=5, callbacks=[wandb.keras.WandbCallback(input_type="images", save_model=False)],
validation_data=(test_images, test_labels))
In [2]:
sweep_config = {
'method': 'grid',
'parameters': {
'hidden_nodes': {
'values': [32, 64, 96, 128, 256]
}
}
}
In [3]:
import wandb
sweep_id = wandb.sweep(sweep_config)
In [ ]:
wandb.agent(sweep_id, function=train)