civ
132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
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])
|