File size: 3,624 Bytes
13d3ba0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
#!/usr/bin/env python3
import sys
import gguf
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers
def train(model_name):
# 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
# Make sure images have shape (28, 28, 1)
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),
layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
layers.MaxPooling2D(pool_size=(2, 2)),
layers.Flatten(),
layers.Dropout(0.5),
layers.Dense(num_classes, activation="softmax"),
]
)
model.summary()
batch_size = 128
epochs = 15
model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)
score = model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", score[0])
print("Test accuracy:", score[1])
model.save(model_name)
print("Keras model saved to '" + model_name + "'")
def convert(model_name):
model = keras.models.load_model(model_name)
gguf_model_name = model_name + ".gguf"
gguf_writer = gguf.GGUFWriter(gguf_model_name, "mnist-cnn")
kernel1 = model.layers[0].weights[0].numpy()
kernel1 = np.moveaxis(kernel1, [2,3], [0,1])
kernel1 = kernel1.astype(np.float16)
gguf_writer.add_tensor("kernel1", kernel1, raw_shape=(32, 1, 3, 3))
bias1 = model.layers[0].weights[1].numpy()
bias1 = np.repeat(bias1, 26*26)
gguf_writer.add_tensor("bias1", bias1, raw_shape=(1, 32, 26, 26))
kernel2 = model.layers[2].weights[0].numpy()
kernel2 = np.moveaxis(kernel2, [0,1,2,3], [2,3,1,0])
kernel2 = kernel2.astype(np.float16)
gguf_writer.add_tensor("kernel2", kernel2, raw_shape=(64, 32, 3, 3))
bias2 = model.layers[2].weights[1].numpy()
bias2 = np.repeat(bias2, 11*11)
gguf_writer.add_tensor("bias2", bias2, raw_shape=(1, 64, 11, 11))
dense_w = model.layers[-1].weights[0].numpy()
dense_w = dense_w.transpose()
gguf_writer.add_tensor("dense_w", dense_w, raw_shape=(10, 1600))
dense_b = model.layers[-1].weights[1].numpy()
gguf_writer.add_tensor("dense_b", dense_b)
gguf_writer.write_header_to_file()
gguf_writer.write_kv_data_to_file()
gguf_writer.write_tensors_to_file()
gguf_writer.close()
print("Model converted and saved to '{}'".format(gguf_model_name))
if __name__ == '__main__':
if len(sys.argv) < 3:
print("Usage: %s <train|convert> <model_name>".format(sys.argv[0]))
sys.exit(1)
if sys.argv[1] == 'train':
train(sys.argv[2])
elif sys.argv[1] == 'convert':
convert(sys.argv[2])
else:
print("Usage: %s <train|convert> <model_name>".format(sys.argv[0]))
sys.exit(1)
|