|
|
from tensorflow import keras |
|
|
|
|
|
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data() |
|
|
|
|
|
x_train = x_train / 255.0 |
|
|
x_test = x_test / 255.0 |
|
|
|
|
|
model = keras.models.Sequential([ |
|
|
keras.layers.Flatten(input_shape=(28, 28)), |
|
|
keras.layers.Dense(512, activation='relu'), |
|
|
keras.layers.Dense(512, activation='relu'), |
|
|
keras.layers.Dense(10, activation='softmax') |
|
|
]) |
|
|
|
|
|
model.compile(optimizer=keras.optimizers.Adam(0.001), |
|
|
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), |
|
|
metrics=[keras.metrics.SparseCategoricalAccuracy()]) |
|
|
|
|
|
model.fit(x_train, y_train, validation_data=(x_test, y_test), epochs=6) |
|
|
|
|
|
model.save_weights('./weights/mnist.weights.h5') |