in tensorflow-train-in-sagemaker-deploy-with-lambda/mnist-2.py [0:0]
def model(x_train, y_train, x_test, y_test):
"""Generate a simple model"""
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(1024, activation=tf.nn.relu),
tf.keras.layers.Dropout(0.4),
tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.fit(x_train, y_train)
model.evaluate(x_test, y_test)
return model