example_mnist.py

import matplotlib.pyplot as plt
from biolayer import BioLinear
from visualization import LayerVisualizer
from torchvision.datasets import MNIST

MNIST_DIR = '~/DataSets/'

num_inputs = 784
batch_size = 32
learning_rate = 2e-3
num_neurons = 5 * 5

M = MNIST(MNIST_DIR, download=True).data.detach().view(-1, 28*28).float()
M = (M - M.mean(0)) / 255.0
M.cuda()

bio_linear = BioLinear(num_inputs, num_neurons, bias=True)
bio_linear.cuda()
vis = LayerVisualizer(bio_linear.weight)

avgs = []
stds = []
try:
	for weight in bio_linear.train(M, batch_size=batch_size, epsilon=learning_rate):
		vis.update()
		avgs.append(weight.mean().item())
		stds.append(weight.std().item())
except KeyboardInterrupt:
	pass

plt.plot(avgs, label="avg")
plt.plot(stds, label="std")
plt.legend()
plt.show()