Multiclass Classifier Accuracy

pytorch-multiclass-accuracy.py

test_correct = 0
train_correct = 0

output = (net(mnist_test_X) > 0.5).float()
for i, v in enumerate(output):
	if v.argmax() == mnist_test_y[i].argmax():
		test_correct += 1

output = (net(mnist_train_X) > 0.5).float()
for i, v in enumerate(output):
	if v.argmax() == mnist_train_y[i].argmax():
		train_correct += 1

print("Train Accuracy: {:.3f}, Test Accuracy: {:.3f}".format(train_correct / mnist_train_y.shape[0], test_correct / mnist_test_y.shape[0]))