mnist-c

main.c (21233B)

#include <ncurses.h>

#include <sys/random.h>
#include <math.h>
#include <err.h>
#include <signal.h>
#include <assert.h>
#include <endian.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>

#define ARRLEN(x) (sizeof(x)/sizeof((x)[0]))

enum {
	U8 = 0x08,
	I8 = 0x09,
	I16 = 0x0B,
	I32 = 0x0C,
	F32 = 0x0D,
	F64 = 0x0E
};

enum {
	IDENTITY,
	SIGMOID,
	SOFTMAX
};

struct idx {
	void *data;
	uint32_t *dim;
	uint8_t dims;
	uint8_t dtype;
};

struct layer_spec {
	int activation;
	bool has_bias;
	size_t len;
};

struct layer {
	int activation;
	bool has_bias;
	size_t len, nodes;
	double *activity;
	double *input;
	double *derivs;
};

struct nn {
	char *filepath;
	struct layer *layer;
	size_t layers;
	struct layer *input, *output;
	/* 3d matrix, [layer][source][target] */
	double ***weights;
	double ***deltas;
};

static const struct layer_spec layers[] = {
	{ IDENTITY, true, 28 * 28 },
	{ SIGMOID, true, 20 },
	{ SOFTMAX, false, 10 },
};

static const uint8_t idx_dtype_size[0x100] = {
	[U8] = 1,
	[I8] = 1,
	[I16] = 2,
	[I32] = 4,
	[F32] = 4,
	[F64] = 8
};

static bool quit = false;

void
train_stop(int sig)
{
	quit = true;
	printf("QUIT\n");
}

double
dbl_be64toh(double d)
{
	uint64_t tmp;

	tmp = *(uint64_t*)&d;
	tmp = be64toh(tmp);
	return *(double*)&tmp;
}

double
dbl_htobe64(double d)
{
	uint64_t tmp;

	tmp = *(uint64_t*)&d;
	tmp = htobe64(tmp);
	return *(double*)&tmp;
}

void
idx_load(struct idx *idx, FILE *file, const char *path)
{
	uint8_t header[2];
	uint32_t count;
	uint32_t *counts;
	size_t size;
	int i;

	if (fread(header, 1, 2, file) != 2)
		errx(1, "Missing idx header (%s)", path);

	if (header[0] || header[1])
		errx(1, "Invalid idx header (%s)", path);

	if (fread(&idx->dtype, 1, 1, file) != 1)
		errx(1, "Missing idx data type (%s)", path);

	if (fread(&idx->dims, 1, 1, file) != 1)
		errx(1, "Missing idx dims (%s)", path);

	if (!idx_dtype_size[idx->dtype])
		errx(1, "Invalid idx data type (%s)", path);

	idx->dim = malloc(idx->dims * sizeof(uint32_t));
	if (!idx->dim) err(1, "malloc");

	size = 1;
	for (i = 0; i < idx->dims; i++) {
		if (fread(&count, 4, 1, file) != 1)
			errx(1, "Missing %i. dimension size (%s)", i + 1, path);
		idx->dim[i] = be32toh(count);
		size *= idx->dim[i];
	}

	idx->data = malloc(size * idx_dtype_size[idx->dtype]);;
	if (!idx->dtype) err(1, "malloc");

	if (fread(idx->data, idx_dtype_size[idx->dtype], size, file) != size)
		errx(1, "Incomplete data section (%s)", path);
}

void
idx_free(struct idx *idx)
{
	free(idx->data);
	idx->data = NULL;
	free(idx->dim);
	idx->dim = NULL;
}

void
idx_load_single(struct idx *idx, const char *path)
{
	FILE *file;

	file = fopen(path, "r");
	if (!file) err(1, "fopen (%s)", path);
	idx_load(idx, file, path);
	fclose(file);
}

void
idx_load_images(struct idx *idx, const char *path)
{
	idx_load_single(idx, path);
	assert(idx->dims == 3);
	assert(idx->dim[1] == 28 && idx->dim[2] == 28);
	assert(idx->dtype == U8);
}

void
idx_load_labels(struct idx *idx, const char *path)
{
	idx_load_single(idx, path);
	assert(idx->dims == 1);
	assert(idx->dtype == U8);
}

void
idx_save(struct idx *idx, FILE *file, const char *path)
{
	uint8_t header[2];
	uint32_t count;
	size_t size;
	int i;

	memset(header, 0, 2);
	if (fwrite(&header, 1, 2, file) != 2)
		err(1, "fwrite (%s)", path);

	if (fwrite(&idx->dtype, 1, 1, file) != 1)
		err(1, "fwrite (%s)", path);

	if (fwrite(&idx->dims, 1, 1, file) != 1)
		err(1, "fwrite (%s)", path);

	size = 1;
	for (i = 0; i < idx->dims; i++) {
		count = htobe32(idx->dim[i]);
		if (fwrite(&count, 4, 1, file) != 1)
			err(1, "fwrite (%s)", path);
		size *= idx->dim[i];
	}

	if (fwrite(idx->data, idx_dtype_size[idx->dtype], size, file) != size)
		err(1, "fwrite (%s)", path);
}

void
nn_init(struct nn *nn, const struct layer_spec *spec, size_t layers)
{
	int l, k;

	nn->filepath = NULL;
	nn->layers = layers;
	nn->layer = malloc(sizeof(struct layer) * nn->layers);

	for (l = 0; l < nn->layers; l++) {
		nn->layer[l].len = spec[l].len;
		nn->layer[l].has_bias = spec[l].has_bias;
		nn->layer[l].activation = spec[l].activation;

		nn->layer[l].nodes = spec[l].len + spec[l].has_bias;

		nn->layer[l].input = calloc(nn->layer[l].nodes, sizeof(double));
		if (!nn->layer[l].input) err(1, "malloc");

		nn->layer[l].activity = calloc(nn->layer[l].nodes, sizeof(double));
		if (!nn->layer[l].activity) err(1, "malloc");

		nn->layer[l].derivs = calloc(nn->layer[l].nodes, sizeof(double));
		if (!nn->layer[l].derivs) err(1, "malloc");
	}

	nn->input = &nn->layer[0];
	nn->output = &nn->layer[nn->layers - 1];

	nn->deltas = malloc((nn->layers - 1) * sizeof(double *));
	if (!nn->deltas) err(1, "malloc");

	nn->weights = malloc((nn->layers - 1) * sizeof(double *));
	if (!nn->weights) err(1, "malloc");

	for (l = 0; l < nn->layers - 1; l++) {
		nn->deltas[l] = malloc(nn->layer[l].nodes * sizeof(double *));
		if (!nn->deltas[l]) err(1, "malloc");
		for (k = 0; k < nn->layer[l].nodes; k++) {
			nn->deltas[l][k] = calloc(nn->layer[l+1].len,
				sizeof(double));
			if (!nn->deltas[l][k]) err(1, "malloc");
		}

		nn->weights[l] = malloc(nn->layer[l].nodes * sizeof(double *));
		if (!nn->weights[l]) err(1, "malloc");
		for (k = 0; k < nn->layer[l].nodes; k++) {
			nn->weights[l][k] = calloc(nn->layer[l+1].len,
				sizeof(double));
			if (!nn->weights[l][k]) err(1, "malloc");
		}
	}
}

void
nn_gen(struct nn *nn)
{
	int l, s, t;
	uint32_t val;

	/* initial weights */
	for (l = 0; l < nn->layers - 1; l++) {
		for (s = 0; s < nn->layer[l].nodes; s++) {
			for (t = 0; t < nn->layer[l+1].len; t++) {
				if (getrandom(&val, 4, 0) != 4)
					err(1, "getrandom");
				nn->weights[l][s][t] =
					((val / (double) 0xFFFFFFFF) - 0.5)
					/ nn->layer[l].nodes;
			}
		}
	}
}

void
nn_load(struct nn *nn, const char *path)
{
	FILE *file;
	struct idx idx;
	double weight;
	int l, s, t;
	int snodes;

	nn->filepath = strdup(path);
	if (!nn->filepath) err(1, "strdup");

	file = fopen(path, "r");
	if (!file) err(1, "fopen (%s)", path);

	/* load weights */
	for (l = 0; l < nn->layers - 1; l++) {
		idx_load(&idx, file, path);
		assert(idx.dtype == F64);
		assert(idx.dims == 2);
		assert(idx.dim[0] == nn->layer[l].nodes);
		assert(idx.dim[1] == nn->layer[l+1].len);
		snodes = nn->layer[l].nodes;
		for (s = 0; s < nn->layer[l].nodes; s++) {
			for (t = 0; t < nn->layer[l+1].len; t++) {
				weight = ((double*)idx.data)[t * snodes + s];
				nn->weights[l][s][t] = dbl_be64toh(weight);
			}
		}
		idx_free(&idx);
	}

	fclose(file);
}

void
nn_save(struct nn *nn, const char *path)
{
	FILE *file;
	struct idx idx;
	double weight;
	int l, s, t;
	int snodes;

	file = fopen(path, "w+");
	if (!file) err(1, "fopen (%s)", path);

	idx.dims = 2;
	idx.dim = malloc(idx.dims * sizeof(uint32_t));
	if (!idx.dim) err(1, "malloc");
	idx.dtype = F64;

	/* save weights */
	for (l = 0; l < nn->layers - 1; l++) {
		idx.data = malloc(nn->layer[l].nodes
			* nn->layer[l+1].len * sizeof(double));
		if (!idx.data) err(1, "malloc");
		snodes = nn->layer[l].nodes;
		for (s = 0; s < nn->layer[l].nodes; s++) {
			for (t = 0; t < nn->layer[l+1].len; t++) {
				weight = dbl_htobe64(nn->weights[l][s][t]);
				((double *)idx.data)[t * snodes + s] = weight;
			}
		}
		idx.dim[0] = nn->layer[l].nodes;
		idx.dim[1] = nn->layer[l+1].len;
		idx_save(&idx, file, path);
		free(idx.data);
	}

	free(idx.dim);
	fclose(file);
}

void
nn_free(struct nn *nn)
{
	int l, k;

	free(nn->filepath);

	for (l = 0; l < nn->layers; l++) {
		free(nn->layer[l].derivs);
		free(nn->layer[l].activity);
		free(nn->layer[l].input);
		if (l < nn->layers - 1) {
			for (k = 0; k < nn->layer[l].nodes; k++) {
				free(nn->deltas[l][k]);
				free(nn->weights[l][k]);
			}
			free(nn->deltas[l]);
			free(nn->weights[l]);
		}
	}

	free(nn->layer);
	free(nn->weights);
	free(nn->deltas);
}

void
nn_fwdprop_layer(struct nn *nn, int l)
{
	struct layer *sl, *tl;
	double expsum, weight, max;
	int s, t;

	sl = &nn->layer[l];
	tl = &nn->layer[l+1];

	if (tl->has_bias)
		tl->activity[tl->len] = 1.0;
	for (t = 0; t < tl->len; t++) {
		tl->input[t] = 0;
		for (s = 0; s < sl->nodes; s++) {
			tl->input[t] += sl->activity[s]
				* nn->weights[l][s][t];
		}
	}

	switch (tl->activation) {
	case IDENTITY:
		for (t = 0; t < tl->len; t++)
			tl->activity[t] = tl->input[t];
		break;
	case SIGMOID:
		for (t = 0; t < tl->len; t++)
			tl->activity[t] = 1 / (1 + exp(-tl->input[t]));
		break;
	case SOFTMAX:
		max = tl->input[0];
		for (t = 0; t < tl->len; t++)
			max = tl->input[t] > max ? tl->input[t] : max;
		expsum = 0;
		for (t = 0; t < tl->len; t++)
			expsum += exp(tl->input[t] - max);
		for (t = 0; t < tl->len; t++)
			tl->activity[t] = exp(tl->input[t] - max) / expsum;
		break;
	default:
		errx(1, "Unknown activation function (%i)", tl->activation);
	};
}


void
nn_fwdprop(struct nn *nn, uint8_t *image)
{
	int i, l;

	nn->layer[0].activity[nn->layer[0].len] = 1.0;
	for (i = 0; i < nn->layer[0].len; i++)
		nn->layer[0].activity[i] = image[i] ? 1.0 : 0.0;

	for (l = 0; l < nn->layers - 1; l++)
		nn_fwdprop_layer(nn, l);
}

void
nn_backprop_layer(struct nn *nn, int l)
{
	struct layer *sl, *tl;
	int s, t, i;
	double sum;

	sl = &nn->layer[l-1];
	tl = &nn->layer[l];

	for (s = 0; s < sl->nodes; s++)
		sl->derivs[s] = 0;

	switch (nn->layer[l].activation) {
	case IDENTITY:
		for (t = 0; t < tl->len; t++) {
			for (s = 0; s < sl->nodes; s++) {
				sl->derivs[s] += tl->derivs[t]
					* nn->weights[l-1][s][t];
			}
		}
		break;
	case SIGMOID:
		for (t = 0; t < tl->len; t++) {
			/* derivative of activation function */
			tl->derivs[t] *= tl->activity[t] * (1 - tl->activity[t]);
			for (s = 0; s < sl->nodes; s++) {
				sl->derivs[s] += tl->derivs[t]
					* nn->weights[l-1][s][t];
			}
		}
		break;
	case SOFTMAX:
		/* derivative of softmax function
		 * (each input i influences activity t) */
		for (t = 0; t < tl->nodes; t++) {
			sum = 0;
			for (i = 0; i < tl->nodes; i++) {
				sum += tl->derivs[i] * tl->activity[i]
					* ((t == i) - tl->activity[t]);
			}
			for (s = 0; s < sl->nodes; s++)
				sl->derivs[s] += sum * nn->weights[l-1][s][t];
		}
		break;
	}
}

void
nn_backprop(struct nn *nn, uint8_t label)
{
	int i, l;

	l = nn->layers - 1;
	for (i = 0; i < nn->layer[l].len; i++) {
		/* derivative of error: 1/2 (label - out)^2 */
		nn->layer[l].derivs[i] = nn->layer[l].activity[i]
			- (label == i ? 1.0 : 0.0);
	}

	/* generate derivs of err / z_i per node */
	for (l = nn->layers - 1; l >= 1; l--)
		nn_backprop_layer(nn, l);
}

void
nn_debug(struct nn *nn)
{
	int l, s, t;

	printf("WEIGHTS:\n");
	for (l = 0; l < nn->layers - 1; l++) {
		printf("LAYER %i\n", l);
		for (s = 0; s < nn->layer[l].nodes; s++) {
			for (t = 0; t < nn->layer[l+1].len; t++) {
				printf("%0.3F ", nn->weights[l][s][t]);
			}
			printf("\n");
		}
		printf("\n");
	}

	printf("DELTAS:\n");
	for (l = 0; l < nn->layers - 1; l++) {
		printf("LAYER %i\n", l);
		for (s = 0; s < nn->layer[l].nodes; s++) {
			for (t = 0; t < nn->layer[l+1].len; t++) {
				printf("%0.3F ", nn->deltas[l][s][t]);
			}
			printf("\n");
		}
		printf("\n");
	}
}

void
nn_debug_prediction(struct nn *nn, uint8_t label)
{
	int k;

	printf("%i : ", label);
	for (k = 0; k < nn->output->len; k++)
		printf("%2.0F ", 100 * nn->output->activity[k]);
	printf("\n");
}

int
weight_color(double weight)
{
	int color;

	if (weight >= 0) {
		if (weight < 0.01) {
			color = 22;
		} else if (weight < 0.1) {
			color = 28;
		} else if (weight < 1) {
			color = 34;
		} else if (weight < 10) {
			color = 40;
		} else {
			color = 46;
		}
	} else {
		if (weight > -0.01) {
			color = 52;
		} else if (weight > -0.1) {
			color = 88;
		} else if (weight > -1) {
			color = 124;
		} else if (weight > -10) {
			color = 160;
		} else {
			color = 196;
		}
	}

	return color;
}

void
nn_dump(struct nn *nn)
{
	int l, s, t, x, y;
	double weight;

	printf("\n");
	for (t = 0; t < nn->layer[1].len; t++) {
		printf("INPUT -> HIDDEN %i\n", t);
		for (y = 0; y < 28; y++) {
			for (x = 0; x < 28 + (y == 27); x++) {
				weight = nn->weights[0][y * 28 + x][t];
				printf("\x1b[38:5:%im%s\x1b[0m",
					weight_color(weight),
					fabs(weight) >= 0.0001 ? "▮" : " ");
			}
			printf("\n");
		}
		printf("\n");
	}

	if (nn->layers > 2) {
		printf("HIDDEN -> OUTPUT\n");
		for (t = 0; t < nn->layer[2].len; t++) {
			for (s = 0; s < nn->layer[1].nodes; s++) {
				weight = nn->weights[1][s][t];
				printf("\x1b[38:5:%im%s\x1b[0m",
					weight_color(weight),
					fabs(weight) >= 0.0001 ? "▮" : " ");
			}
			printf("\n");
		}
	}

}

void
nn_reset_deltas(struct nn *nn)
{
	int l, s, t;

	for (l = 0; l < nn->layers - 1; l++) {
		for (s = 0; s < nn->layer[l].nodes; s++) {
			for (t = 0; t < nn->layer[l+1].len; t++)
				nn->deltas[l][s][t] = 0;
		}
	}
}

void
nn_update_deltas(struct nn *nn, double learning_rate)
{
	int l, s, t;
	double gradw;

	/* generate deltas for weights from err / z_i */
	for (l = nn->layers - 1; l >= 1; l--) {
		for (t = 0; t < nn->layer[l].len; t++) {
			for (s = 0; s < nn->layer[l-1].nodes; s++) {
				gradw = - nn->layer[l].derivs[t]
					* nn->layer[l-1].activity[s];
				nn->deltas[l-1][s][t] += gradw * learning_rate;
			}
		}
	}
}

void
nn_apply_deltas(struct nn *nn, size_t size)
{
	int l, s, t;

	for (l = nn->layers - 1; l >= 1; l--) {
		for (t = 0; t < nn->layer[l].len; t++) {
			for (s = 0; s < nn->layer[l-1].nodes; s++) {
				nn->weights[l-1][s][t] +=
					nn->deltas[l-1][s][t] / size;
				assert(!isnan(nn->weights[l-1][s][t]));
			}
		}
	}
}

int
nn_result(struct nn *nn)
{
	double max;
	int k, maxi;

	maxi = -1;
	for (k = 0; k < nn->output->len; k++) {
		if (maxi < 0 || nn->output->activity[k] > max) {
			max = nn->output->activity[k];
			maxi = k;
		}
	}

	return maxi;
}

double
nn_test(struct nn *nn)
{
	struct idx images;
	struct idx labels;
	size_t hits, total;
	int i, k, res;
	uint8_t label;
	double max;

	idx_load_images(&images, "data/test-images.idx");
	idx_load_labels(&labels, "data/test-labels.idx");

	total = hits = 0;
	for (i = 0; i < images.dim[0]; i++) {
		nn_fwdprop(nn, images.data + i * nn->input->len);
		label = *(uint8_t*)(labels.data + i);
		nn_debug_prediction(nn, label);
		res = nn_result(nn);
		if (res == label) hits++;
		total++;
	}

	idx_free(&images);
	idx_free(&labels);

	return 1.F * hits / total;
}

double
nn_batch(struct nn *nn, struct idx *images, struct idx *labels,
	size_t batch_size, double learning_rate)
{
	double lerror, error;
	uint32_t idx;
	uint8_t label;
	size_t i, k;

	nn_reset_deltas(nn);

	error = 0;
	for (i = 0; i < batch_size; i++) {
		if (getrandom(&idx, 4, 0) != 4)
			err(1, "getrandom");
		idx = idx % images->dim[0];
		nn_fwdprop(nn, images->data + idx * nn->layer[0].len);
		label = *(uint8_t *)(labels->data + idx);
		for (k = 0; k < nn->output->nodes; k++) {
			lerror = nn->output->activity[k]
				- (k == label ? 1.0 : 0.0);
			error += 0.5 * lerror * lerror;
		}
		//nn_debug_prediction(nn, label);
		nn_backprop(nn, label);
		nn_update_deltas(nn, learning_rate);
	}

	nn_apply_deltas(nn, batch_size);

	return error / batch_size;
}

void
nn_train(struct nn *nn, size_t epochs,
	size_t batch_size, double learning_rate)
{
	struct idx images;
	struct idx labels;
	double error;
	int epoch, i;

	signal(SIGINT, train_stop);

	idx_load_images(&images, "data/train-images.idx");
	idx_load_labels(&labels, "data/train-labels.idx");

	for (epoch = 0; epoch < epochs; epoch++) {
		/* TODO: ensure using all images in one epoch */
		for (i = 0; i < images.dim[0] / batch_size; i++) {
			error = nn_batch(nn, &images, &labels,
				batch_size, learning_rate);
			if (i % 100 == 0) {
				//nn_debug(nn);
				//nn_dump(nn);
				printf("Batch %i / %lu => %2.5F\n", i + 1,
					images.dim[0] / batch_size, error);
			}
			if (quit) {
				nn_save(nn, nn->filepath);
				exit(1);
			}
		}
	}

	idx_free(&images);
	idx_free(&labels);
}

void
nn_trainvis(struct nn *nn, size_t batch_size, double learning_rate)
{
	struct idx images;
	struct idx labels;
	double error, weight;
	int epoch, i;
	int t, x, y;
	int sx, sy;
	bool show;

	/* display weights visually after each batch
	 * and adjust batch frequency via UP / DOWN */
	signal(SIGINT, train_stop);

	idx_load_images(&images, "data/train-images.idx");
	idx_load_labels(&labels, "data/train-labels.idx");

	printf("\x1b[?25l"); /* hide cursor */
	printf("\x1b[2J"); /* clear screen */

	while (!quit) {
		error = nn_batch(nn, &images, &labels,
			batch_size, learning_rate);
		if (quit) {
			nn_save(nn, nn->filepath);
			break;
		}

		printf("\x1b[%i;%iHTraining error: %F", 2, 0, error);

		assert(nn->layers > 1);
		for (t = 0; t < nn->layer[1].len; t++) {
			sy = (t >= nn->layer[1].len / 2) ? 35 : 5;
			sx = 2 + 30 * (t % (nn->layer[1].len / 2));
			for (y = 0; y < 28; y++) {
				for (x = 0; x < 28 + (y == 27); x++) {
					weight = nn->weights[0][y * 28 + x][t];
					show = fabs(weight) >= 0.0001;
					printf("\x1b[%i;%iH", sy + y, sx + x);
					printf("\x1b[38:5:%im%s\x1b[0m",
						weight_color(weight),
						show ? "▮" : " ");
				}
			}
		}
	}

	printf("\x1b[?25h"); /* show cursor */
	printf("\x1b[2J"); /* clear screen */ 

	idx_free(&images);
	idx_free(&labels);
}

void
nn_predict(struct nn *nn)
{
	uint8_t image[28*28];
	WINDOW *win;
	MEVENT event;
	int width, height;
	int startx, starty;
	int x, y, c, i, label;
	bool evaluate;

	/* gui interface to draw input and show prediction */

	/* TODO: 256 color support, is this portable? */
	setenv("TERM", "xterm-1002", 1);

	initscr();
	keypad(stdscr, true);
	noecho();
	cbreak();
	curs_set(0);

	mousemask(ALL_MOUSE_EVENTS | REPORT_MOUSE_POSITION, NULL);

	win = NULL;
	label = -1;
	evaluate = true;
	memset(image, 0, sizeof(image));
	while (!quit) {
		width = getmaxx(stdscr);
		height = getmaxy(stdscr);
		assert(width >= 30 && height >= 31);

		startx = (width - 30) / 2;
		starty = (height - 30) / 2;

		if (evaluate) {
			nn_fwdprop(nn, image);
			label = nn_result(nn);
			evaluate = false;
		}

		if (!win) {
			win = newwin(30, 30, starty, startx);
			if (!win) err(1, "newwin");
		} else {
			mvwin(win, starty, startx);
		}

		clear();

		mvprintw(starty - 1, startx - 1, "Predictions: ");
		for (i = 0; i < nn->output->len; i++) {
			if (i == label) attron(A_UNDERLINE);
			if (nn->output->activity[i] >= 0.2)
				attron(A_BOLD);
			printw("%i", i);
			attroff(A_BOLD);
			attroff(A_UNDERLINE);
			printw(" ");
		}
		refresh();

		box(win, 0, 0);
		for (y = 0; y < 28; y++) {
			for (x = 0; x < 28; x++) {
				if (image[y * 28 + x])
					mvwaddch(win, 1 + y, 1 + x, ACS_BLOCK);
				else
					mvwaddch(win, 1 + y, 1 + x, ' ');
			}
		}
		wrefresh(win);

		switch ((c = getch())) {
		case KEY_MOUSE:
			if (getmouse(&event) != OK)
				err(1, "getmouse");
			x = event.x - (startx + 1);
			y = event.y - (starty + 1);
			if (x < 0 || x >= 28) continue;
			if (y < 0 || y >= 28) continue;
			image[y * 28 + x] = 1;
			if (y > 0) image[(y-1) * 28 + x] = 1;
			if (y < 27) image[(y+1) * 28 + x] = 1;
			if (x > 0) image[y * 28 + x - 1] = 1;
			if (x < 27) image[y * 28 + x + 1] = 1;
			if (event.bstate & BUTTON1_RELEASED ||
					event.bstate & BUTTON2_RELEASED)
				evaluate = true;
			break;
		case 'c':
			memset(image, 0, sizeof(image));
			break;
		case 'q':
			quit = true;
			break;
		}
	}

	delwin(win);
	endwin();
}

void
dump_sample(const char *set, size_t index)
{
	struct idx images, labels;
	uint8_t pix, *image;
	int x, y;

	if (!strcmp(set, "train")) {
		idx_load_images(&images, "data/train-images.idx");
		idx_load_labels(&labels, "data/train-labels.idx");
	} else if (!strcmp(set, "test")) {
		idx_load_images(&images, "data/test-images.idx");
		idx_load_labels(&labels, "data/test-labels.idx");
	} else {
		errx(1, "Unknown dataset (%s)", set);
	}

	image = images.data + 28 * 28 * index;
	assert(index < images.dim[0]);
	for (y = 0; y < 28; y++) {
		for (x = 0; x < 28; x++) {
			pix = image[y * 28 + x];
			printf("%s", pix ? "▮" : " ");
		}
		printf("\n");
	}

	printf("Label: %i\n", *(uint8_t *)(labels.data + index));

	idx_free(&images);
}

int
main(int argc, const char **argv)
{
	struct nn nn;

	if (argc == 2 && !strcmp(argv[1], "gen")) {
		nn_init(&nn, layers, ARRLEN(layers));
		nn_gen(&nn);
		nn_save(&nn, ".nn");
		nn_free(&nn);
	} else if (argc == 2 && !strcmp(argv[1], "train")) {
		nn_init(&nn, layers, ARRLEN(layers));
		nn_load(&nn, ".nn");
		nn_train(&nn, 10, 5, 0.005);
		nn_save(&nn, ".nn");
		nn_free(&nn);
	} else if (argc == 2 && !strcmp(argv[1], "trainvis")) {
		nn_init(&nn, layers, ARRLEN(layers));
		nn_load(&nn, ".nn");
		nn_trainvis(&nn, 5, 0.005);
		nn_save(&nn, ".nn");
		nn_free(&nn);
	} else if (argc == 2 && !strcmp(argv[1], "predict")) {
		nn_init(&nn, layers, ARRLEN(layers));
		nn_load(&nn, ".nn");
		nn_predict(&nn);
		nn_free(&nn);
	} else if (argc == 2 && !strcmp(argv[1], "test")) {
		nn_init(&nn, layers, ARRLEN(layers));
		nn_load(&nn, ".nn");
		printf("Accuracy: %F\n", nn_test(&nn));
		nn_free(&nn);
	} else if (argc == 2 && !strcmp(argv[1], "dump")) {
		nn_init(&nn, layers, ARRLEN(layers));
		nn_load(&nn, ".nn");
		nn_dump(&nn);
		nn_free(&nn);
	} else if (argc == 4 && !strcmp(argv[1], "sample")) {
		dump_sample(argv[2], atoi(argv[3]));
	} else {
		printf("Commands: gen train trainvis predict test dump sample\n");
	}
}