make CIFAR-10 nin model and scripts run on PyTorch 1.7.0

CIFAR_10/main.py

@@ -10,6 +10,8 @@
 import util
 import torch.nn as nn
 import torch.optim as optim
+import torchvision
+from torchvision import transforms
 
 from models import nin
 from torch.autograd import Variable
@@ -29,11 +31,16 @@ def save_state(model, best_acc):
 def train(epoch):
     model.train()
     for batch_idx, (data, target) in enumerate(trainloader):
+        if torch.cuda.is_available():
+            data = data.cuda()
+            target = target.cuda()
+
         # process the weights including binarization
         bin_op.binarization()
 
         # forwarding
-        data, target = Variable(data.cuda()), Variable(target.cuda())
+        data, target = Variable(data), Variable(target)
+
         optimizer.zero_grad()
         output = model(data)
 
@@ -60,7 +67,9 @@ def test():
     correct = 0
     bin_op.binarization()
     for data, target in testloader:
-        data, target = Variable(data.cuda()), Variable(target.cuda())
+        if torch.cuda.is_available():
+            data, target = data.cuda(), target.cuda()
+        data, target = Variable(data), Variable(target)
 
         output = model(data)
         test_loss += criterion(output, target).data.item()
@@ -90,8 +99,6 @@ def adjust_learning_rate(optimizer, epoch):
 if __name__=='__main__':
     # prepare the options
     parser = argparse.ArgumentParser()
-    parser.add_argument('--cpu', action='store_true',
-            help='set if only CPU is available')
     parser.add_argument('--data', action='store', default='./data/',
             help='dataset path')
     parser.add_argument('--arch', action='store', default='nin',
@@ -110,16 +117,20 @@ def adjust_learning_rate(optimizer, epoch):
     torch.cuda.manual_seed(1)
 
     # prepare the data
-    if not os.path.isfile(args.data+'/train_data'):
+    if not os.path.isdir(args.data):
         # check the data path
         raise Exception\
                 ('Please assign the correct data path with --data <DATA_PATH>')
+
+    to_tensor_transformer = transforms.Compose([
+        transforms.ToTensor(),
+        ])
+    trainset = torchvision.datasets.CIFAR10(args.data, train=True, download=True, transform=to_tensor_transformer)
 
-    trainset = data.dataset(root=args.data, train=True)
     trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
             shuffle=True, num_workers=2)
 
-    testset = data.dataset(root=args.data, train=False)
+    testset = torchvision.datasets.CIFAR10(args.data, train=False, download=True, transform=to_tensor_transformer)
     testloader = torch.utils.data.DataLoader(testset, batch_size=100,
             shuffle=False, num_workers=2)
 
@@ -148,7 +159,7 @@ def adjust_learning_rate(optimizer, epoch):
         best_acc = pretrained_model['best_acc']
         model.load_state_dict(pretrained_model['state_dict'])
 
-    if not args.cpu:
+    if torch.cuda.is_available():
         model.cuda()
         model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
     print(model)

CIFAR_10/models/nin.py

@@ -6,13 +6,15 @@ class BinActive(torch.autograd.Function):
     '''
     Binarize the input activations and calculate the mean across channel dimension.
     '''
+    @staticmethod
     def forward(self, input):
         self.save_for_backward(input)
         size = input.size()
         mean = torch.mean(input.abs(), 1, keepdim=True)
         input = input.sign()
         return input, mean
 
+    @staticmethod
     def backward(self, grad_output, grad_output_mean):
         input, = self.saved_tensors
         grad_input = grad_output.clone()
@@ -40,7 +42,7 @@ def __init__(self, input_channels, output_channels,
 
     def forward(self, x):
         x = self.bn(x)
-        x, mean = BinActive()(x)
+        x, mean = BinActive.apply(x)
         if self.dropout_ratio!=0:
             x = self.dropout(x)
         x = self.conv(x)

CIFAR_10/util.py

@@ -35,7 +35,6 @@ def binarization(self):
 
     def meancenterConvParams(self):
         for index in range(self.num_of_params):
-            s = self.target_modules[index].data.size()
             negMean = self.target_modules[index].data.mean(1, keepdim=True).\
                     mul(-1).expand_as(self.target_modules[index].data)
             self.target_modules[index].data = self.target_modules[index].data.add(negMean)