initial release

ggml/examples/mnist/mnist-train-fc.py

@@ -0,0 +1,131 @@
+import gguf
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.datasets as dsets
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+
+import sys
+from time import time
+
+input_size  = 784  # img_size = (28,28) ---> 28*28=784 in total
+hidden_size = 500  # number of nodes at hidden layer
+num_classes = 10   # number of output classes discrete range [0,9]
+num_epochs  = 30   # number of times which the entire dataset is passed throughout the model
+batch_size  = 1000 # the size of input data used for one iteration
+lr          = 1e-3 # size of step
+
+
+class Net(nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super(Net, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_size, num_classes)
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.relu(out)
+        out = self.fc2(out)
+        return out
+
+
+def train(model_path):
+    train_data = dsets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
+    test_data  = dsets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
+
+    assert len(train_data) == 60000
+    assert len(test_data)  == 10000
+
+    kwargs_train_test = dict(batch_size=batch_size, num_workers=4, pin_memory=True)
+    train_gen = torch.utils.data.DataLoader(dataset=train_data, shuffle=True,  **kwargs_train_test)
+    test_gen  = torch.utils.data.DataLoader(dataset=test_data,  shuffle=False, **kwargs_train_test)
+
+    net = Net(input_size, hidden_size, num_classes)
+
+    if torch.cuda.is_available():
+        net.cuda()
+
+    loss_function = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
+
+    t_start = time()
+    for epoch in range(num_epochs):
+        loss_history = []
+        ncorrect = 0
+
+        for i, (images, labels) in enumerate(train_gen):
+            images = Variable(images.view(-1, 28*28))
+            labels = Variable(labels)
+
+            if torch.cuda.is_available():
+                images = images.cuda()
+                labels = labels.cuda()
+
+            optimizer.zero_grad()
+            outputs = net(images)
+            loss = loss_function(outputs, labels)
+
+            loss_history.append(loss.cpu().data)
+            _, predictions = torch.max(outputs, 1)
+            ncorrect += (predictions == labels).sum()
+
+            loss.backward()
+            optimizer.step()
+
+            if (i + 1)*batch_size % 10000 == 0:
+                loss_mean = np.mean(loss_history)
+                accuracy = ncorrect / ((i + 1) * batch_size)
+                print(
+                    f"Epoch [{epoch+1:02d}/{num_epochs}], "
+                    f"Step [{(i+1)*batch_size:05d}/{len(train_data)}], "
+                    f"Loss: {loss_mean:.4f}, Accuracy: {100*accuracy:.2f}%")
+    print()
+    print(f"Training took {time()-t_start:.2f}s")
+
+    loss_history = []
+    ncorrect = 0
+
+    for i, (images, labels) in enumerate(test_gen):
+        images = Variable(images.view(-1, 28*28))
+        labels = Variable(labels)
+
+        if torch.cuda.is_available():
+            images = images.cuda()
+            labels = labels.cuda()
+
+        outputs = net(images)
+        loss = loss_function(outputs, labels)
+
+        loss_history.append(loss.cpu().data)
+        _, predictions = torch.max(outputs, 1)
+        ncorrect += (predictions == labels).sum().cpu().numpy()
+
+    loss_mean            = np.mean(loss_history)
+    loss_uncertainty     = np.std(loss_history) / np.sqrt(len(loss_history) - 1)
+    accuracy_mean        = ncorrect / (len(test_gen) * batch_size)
+    accuracy_uncertainty = np.sqrt(accuracy_mean * (1.0 - accuracy_mean) / (len(test_gen) * batch_size))
+    print()
+    print(f"Test loss: {loss_mean:.6f}+-{loss_uncertainty:.6f}, Test accuracy: {100*accuracy_mean:.2f}+-{100*accuracy_uncertainty:.2f}%")
+
+    gguf_writer = gguf.GGUFWriter(model_path, "mnist-fc")
+
+    print()
+    print(f"Model tensors saved to {model_path}:")
+    for tensor_name in net.state_dict().keys():
+        data = net.state_dict()[tensor_name].squeeze().cpu().numpy()
+        print(tensor_name, "\t", data.shape)
+        gguf_writer.add_tensor(tensor_name, data)
+
+    gguf_writer.write_header_to_file()
+    gguf_writer.write_kv_data_to_file()
+    gguf_writer.write_tensors_to_file()
+    gguf_writer.close()
+
+
+if __name__ == '__main__':
+    if len(sys.argv) != 2:
+        print(f"Usage: {sys.argv[0]} <model_path>")
+        sys.exit(1)
+    train(sys.argv[1])