civ
92 lines
3.2 KiB
Python
Executable File
92 lines
3.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
from time import time
|
|
import gguf
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
from tensorflow import keras
|
|
from tensorflow.keras import layers
|
|
|
|
|
|
def train(model_path):
|
|
# Model / data parameters
|
|
num_classes = 10
|
|
input_shape = (28, 28, 1)
|
|
|
|
# Load the data and split it between train and test sets
|
|
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
|
|
|
|
# Scale images to the [0, 1] range
|
|
x_train = x_train.astype("float32") / 255
|
|
x_test = x_test.astype("float32") / 255
|
|
x_train = np.expand_dims(x_train, -1)
|
|
x_test = np.expand_dims(x_test, -1)
|
|
print("x_train shape:", x_train.shape)
|
|
print(x_train.shape[0], "train samples")
|
|
print(x_test.shape[0], "test samples")
|
|
|
|
# convert class vectors to binary class matrices
|
|
y_train = keras.utils.to_categorical(y_train, num_classes)
|
|
y_test = keras.utils.to_categorical(y_test, num_classes)
|
|
|
|
model = keras.Sequential(
|
|
[
|
|
keras.Input(shape=input_shape, dtype=tf.float32),
|
|
layers.Conv2D(8, kernel_size=(3, 3), padding="same", activation="relu", dtype=tf.float32),
|
|
layers.MaxPooling2D(pool_size=(2, 2)),
|
|
layers.Conv2D(16, kernel_size=(3, 3), padding="same", activation="relu", dtype=tf.float32),
|
|
layers.MaxPooling2D(pool_size=(2, 2)),
|
|
layers.Flatten(),
|
|
layers.Dense(num_classes, activation="softmax", dtype=tf.float32),
|
|
]
|
|
)
|
|
|
|
model.summary()
|
|
batch_size = 1000
|
|
epochs = 30
|
|
model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
|
|
|
|
t_start = time()
|
|
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)
|
|
print(f"Training took {time()-t_start:.2f}s")
|
|
|
|
score = model.evaluate(x_test, y_test, verbose=0)
|
|
print(f"Test loss: {score[0]:.6f}")
|
|
print(f"Test accuracy: {100*score[1]:.2f}%")
|
|
|
|
gguf_writer = gguf.GGUFWriter(model_path, "mnist-cnn")
|
|
|
|
conv1_kernel = model.layers[0].weights[0].numpy()
|
|
conv1_kernel = np.moveaxis(conv1_kernel, [2, 3], [0, 1])
|
|
gguf_writer.add_tensor("conv1.kernel", conv1_kernel, raw_shape=(8, 1, 3, 3))
|
|
|
|
conv1_bias = model.layers[0].weights[1].numpy()
|
|
gguf_writer.add_tensor("conv1.bias", conv1_bias, raw_shape=(1, 8, 1, 1))
|
|
|
|
conv2_kernel = model.layers[2].weights[0].numpy()
|
|
conv2_kernel = np.moveaxis(conv2_kernel, [0, 1, 2, 3], [2, 3, 1, 0])
|
|
gguf_writer.add_tensor("conv2.kernel", conv2_kernel, raw_shape=(16, 8, 3, 3))
|
|
|
|
conv2_bias = model.layers[2].weights[1].numpy()
|
|
gguf_writer.add_tensor("conv2.bias", conv2_bias, raw_shape=(1, 16, 1, 1))
|
|
|
|
dense_weight = model.layers[-1].weights[0].numpy()
|
|
dense_weight = dense_weight.transpose()
|
|
gguf_writer.add_tensor("dense.weight", dense_weight, raw_shape=(10, 7*7*16))
|
|
|
|
dense_bias = model.layers[-1].weights[1].numpy()
|
|
gguf_writer.add_tensor("dense.bias", dense_bias)
|
|
|
|
gguf_writer.write_header_to_file()
|
|
gguf_writer.write_kv_data_to_file()
|
|
gguf_writer.write_tensors_to_file()
|
|
gguf_writer.close()
|
|
print(f"GGUF model saved to '{model_path}'")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) != 2:
|
|
print(f"Usage: {sys.argv[0]} <model_path>")
|
|
sys.exit(1)
|
|
train(sys.argv[1])
|