keras grid job

main.py

import uuid
import json
import random

import keras
import numpy as np
import tensorflow as tf
import click


def build_model(blob):
    inputs = keras.Input(shape=(blob['num_columns'],), name='img')
    x = keras.layers.Dense(blob['num_columns'], activation='sigmoid')(inputs)
    for i in range(blob['num_layers']-1):
        x = keras.layers.Dense(blob['num_columns'], activation='sigmoid')(x)
    outputs = keras.layers.Dense(blob['num_columns'], activation='sigmoid')(x)
    model = keras.Model(inputs=inputs, outputs=outputs, name='mnist_model')
    model.compile(loss='binary_crossentropy', optimizer=keras.optimizers.Adam(learning_rate=0.1))
    return model


def train_model(mod, blob):
    np.random.seed(42)
    train = np.random.randint(0, 2, (blob['rows'], blob['num_columns']))
    mod.compile(loss=blob['loss_func'], optimizer=keras.optimizers.Adam(learning_rate=0.1))
    return mod.fit(train, train,
                   validation_split=0.2,
                   epochs=blob['epochs'],
                   verbose=0,
                   callbacks=[keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5,
                                                                patience=15, verbose=2,
                                                                mode='auto', min_delta=0.01)])


def build_train(blob):
    tf.set_random_seed(blob['tf_seed'])
    model = build_model(blob)
    hist = train_model(model, blob)
    blob.update(hist.history)
    blob['lr'] = [float(i) for i in blob['lr']]
    return blob


if __name__ == "__main__":
    res = build_train({'num_columns': random.choice([10, 11, 12]), 'num_layers': random.choice([1, 3, 2, 3, 4, 4, 5, 5]),
                       'loss_func': 'binary_crossentropy', 'tf_seed': random.randint(1, 4200),
                       'epochs': 300, 'rows': random.choice([2000, 3000, 4000])})
    with open(f"/Users/vincent/Development/grid/logs/{str(uuid.uuid4())[:14]}.jsonl", "w") as f:
        json.dump(res, f)