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)


Create sweep with ID: 3qwdp8z9
Sweep URL: https://app.wandb.ai/qualcomm/sweeps-sep26/sweeps/3qwdp8z9

In [ ]:
wandb.agent(sweep_id, function=train)


wandb: Agent Starting Run: fwjp479r with config:
	hidden_nodes: 32
wandb: Agent Started Run: fwjp479r