in china/2020_GCR_Kubeflow_Workshop/resources/mnist-tensorflow-jupyter.py [0:0]
def train(train_images, train_labels, epochs, model_summary_path):
if model_summary_path:
logdir=model_summary_path # + datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)
model = keras.Sequential([
keras.layers.Conv2D(input_shape=(28,28,1), filters=8, kernel_size=3,
strides=2, activation='relu', name='Conv1'),
keras.layers.Flatten(),
keras.layers.Dense(10, activation=tf.nn.softmax, name='Softmax')
])
model.summary()
model.compile(optimizer=tf.train.AdamOptimizer(),
loss='sparse_categorical_crossentropy',
metrics=['accuracy']
)
if model_summary_path:
model.fit(train_images, train_labels, epochs=epochs, callbacks=[tensorboard_callback])
else:
model.fit(train_images, train_labels, epochs=epochs)
return model