diff --git a/CIFAR_10/data.py b/CIFAR_10/data.py deleted file mode 100644 index 0cff69d..0000000 --- a/CIFAR_10/data.py +++ /dev/null @@ -1,38 +0,0 @@ -import os -import torch -import cPickle as pickle -import numpy -import torchvision.transforms as transforms - -class dataset(): - def __init__(self, root=None, train=True): - self.root = root - self.train = train - self.transform = transforms.ToTensor() - if self.train: - train_data_path = os.path.join(root, 'train_data') - train_labels_path = os.path.join(root, 'train_labels') - self.train_data = numpy.load(open(train_data_path, 'r')) - self.train_data = torch.from_numpy(self.train_data.astype('float32')) - self.train_labels = numpy.load(open(train_labels_path, 'r')).astype('int') - else: - test_data_path = os.path.join(root, 'test_data') - test_labels_path = os.path.join(root, 'test_labels') - self.test_data = numpy.load(open(test_data_path, 'r')) - self.test_data = torch.from_numpy(self.test_data.astype('float32')) - self.test_labels = numpy.load(open(test_labels_path, 'r')).astype('int') - - def __len__(self): - if self.train: - return len(self.train_data) - else: - return len(self.test_data) - - def __getitem__(self, index): - if self.train: - img, target = self.train_data[index], self.train_labels[index] - else: - img, target = self.test_data[index], self.test_labels[index] - - - return img, target diff --git a/CIFAR_10/main.py b/CIFAR_10/main.py index b273e98..c2ee149 100644 --- a/CIFAR_10/main.py +++ b/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 ') + + 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) diff --git a/CIFAR_10/models/nin.py b/CIFAR_10/models/nin.py index 289f1ee..cbf7f45 100644 --- a/CIFAR_10/models/nin.py +++ b/CIFAR_10/models/nin.py @@ -6,6 +6,7 @@ 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() @@ -13,6 +14,7 @@ def forward(self, input): 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) diff --git a/CIFAR_10/util.py b/CIFAR_10/util.py index feb29bc..4fdaea0 100644 --- a/CIFAR_10/util.py +++ b/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)