Enable GPU device selection. Update trainer.

Author: danielhavir

main.py

@@ -11,7 +11,7 @@
 parser = argparse.ArgumentParser()
 
 # MNIST or CIFAR?
-parser.add_argument('dataset', nargs='?', type=str, default='MNIST', help="'MNIST' or 'CIFAR'")
+parser.add_argument('dataset', nargs='?', type=str, default='MNIST', help="'MNIST' or 'CIFAR' (case insensitive).")
 # Batch size
 parser.add_argument('-bs', '--batch_size', type=int, default=128, help='Batch size.')
 # Epochs
@@ -24,10 +24,15 @@
 parser.add_argument('--lr_decay', type=float, default=0.96, help='Exponential learning rate decay.')
 # Use multiple GPUs?
 parser.add_argument('--multi_gpu', action='store_true', help='Flag whether to use multiple GPUs.')
+# Select GPU device
+parser.add_argument('--gpu_device', type=int, default=None, help='ID of a GPU to use when multiple GPUs are available.')
 # Data directory
-parser.add_argument('--data_path', type=str, default=DATA_PATH, help='Flag whether to use multiple GPUs.')
+parser.add_argument('--data_path', type=str, default=DATA_PATH, help='Path to the MNIST or CIFAR dataset. Alternatively you can set the path as an environmental variable $data.')
 args = parser.parse_args()
 
+if args.gpu_device is not None:
+    torch.cuda.set_device(args.gpu_device)
+
 if args.multi_gpu:
     args.batch_size *= torch.cuda.device_count()
 

trainer.py

@@ -7,6 +7,7 @@
 import torch.optim as optim
 from torch.autograd import Variable
 import os
+from numpy import prod
 from datetime import datetime
 from model import CapsuleNetwork
 from loss import CapsuleLoss
@@ -39,12 +40,17 @@ def __init__(self, loaders, batch_size, learning_rate, num_routing=3, lr_decay=0
 		self.optimizer = optim.Adam(self.net.parameters(), lr=learning_rate)
 		self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=lr_decay)
 		print(8*'#', 'PyTorch Model built'.upper(), 8*'#')
+		print('Num params:', sum([prod(p.size()) for p in self.net.parameters()]))
 
 	def __repr__(self):
 		return repr(self.net)
 
 	def run(self, epochs, classes):
 		print(8*'#', 'Run started'.upper(), 8*'#')
+		eye = torch.eye(len(classes))
+		if self.use_gpu:
+			eye = eye.cuda()
+		
 		for epoch in range(1, epochs+1):
 			for phase in ['train', 'test']:
 				print(f'{phase}ing...'.capitalize())
@@ -58,10 +64,10 @@ def run(self, epochs, classes):
 				correct = 0; total = 0
 				for i, (images, labels) in enumerate(self.loaders[phase]):
 					t1 = time()
-					# One-hot encode labels
-					labels = torch.eye(len(classes))[labels]
 					if self.use_gpu:
 						images, labels = images.cuda(), labels.cuda()
+					# One-hot encode labels
+					labels = eye[labels]
 
 					images, labels = Variable(images), Variable(labels)
 
@@ -95,8 +101,8 @@ def run(self, epochs, classes):
 		error_rate = round((1-accuracy)*100, 2)
 		torch.save(self.net, os.path.join(SAVE_MODEL_PATH, f'{error_rate}_{now}.pth.tar'))
 
-		class_correct = list(0. for i in range(10))
-		class_total = list(0. for i in range(10))
+		class_correct = list(0. for _ in classes)
+		class_total = list(0. for _ in classes)
 		for images, labels in self.loaders['test']:
 			if self.use_gpu:
 				images, labels = images.cuda(), labels.cuda()
@@ -110,6 +116,6 @@ def run(self, epochs, classes):
 				class_total[label] += 1
 
 
-		for i in range(10):
+		for i in range(len(classes)):
 			print('Accuracy of %5s : %2d %%' % (
 				classes[i], 100 * class_correct[i] / class_total[i]))