init

Author: zhirongw

.gitignore

@@ -0,0 +1,10 @@
+data/*
+checkpoint/*
+logs/*
+others/*
+
+*.pyc
+*.bak
+*.log
+*.tar
+*.pth

README.md

@@ -0,0 +1,66 @@
+## Improving Generalization via Scalable Neighborhood Component Analysis
+
+This repo constains the pytorch implementation for the ECCV2018 paper [(arxiv)](https://arxiv.org/pdf/.pdf).
+The project is about deep learning feature representations optimized for
+nearest neighbor classifiers, which may generalize to new object categories.
+
+Much of code is borrowed from the previous [unsupervised learning project](https://arxiv.org/pdf/1805.01978.pdf).
+Please refer to [this repo](https://github.com/zhirongw/lemniscate.pytorch) for more details.
+
+
+## Pretrained Model
+
+Currently, we provide 3 pretrained ResNet models.
+Each release contains the feature representation of all ImageNet training images (600 mb) and model weights (100-200mb).
+You can also get these representations by forwarding the network for the entire ImageNet images.
+
+- [ResNet 18](http://zhirongw.westus2.cloudapp.azure.com/models/snca_resnet18.pth.tar) (top 1 accuracy 70.59%)
+- [ResNet 34](http://zhirongw.westus2.cloudapp.azure.com/models/snca_resnet34.pth.tar) (top 1 accuracy 74.41%)
+- [ResNet 50](http://zhirongw.westus2.cloudapp.azure.com/models/snca_resnet50.pth.tar) (top 1 accuracy 76.57%)
+
+## Nearest Neighbor
+
+Please follow [this link](http://zhirongw.westus2.cloudapp.azure.com/nn.html) for a list of nearest neighbors on ImageNet.
+Results are visualized from our ResNet50 feature, compared with baseline ResNet50 feature, raw image features and supervised features.
+First column is the query image, followed by 20 retrievals ranked by the similarity.
+
+## Usage
+
+Our code extends the pytorch implementation of imagenet classification in [official pytorch release](https://github.com/pytorch/examples/tree/master/imagenet). 
+Please refer to the official repo for details of data preparation and hardware configurations.
+
+- install python2 and [pytorch=0.3](http://pytorch.org)
+
+- clone this repo: `git clone https://github.com/zhirongw/snca.pytorch`
+
+- Training on ImageNet:
+
+  `python main.py DATAPATH --arch resnet18 -j 32 --temperature 0.05 --low-dim 128 -b 256 `
+
+  - During training, we monitor the supervised validation accuracy by K nearest neighbor with k=1, as it's faster, and gives a good estimation of the feature quality.
+
+- Testing on ImageNet:
+
+  `python main.py DATAPATH --arch resnet18 --resume input_model.pth.tar -e` runs testing with default K=30 neighbors.
+
+- Training on CIFAR10:
+
+  `python cifar.py --nce-t 0.05 --lr 0.1`
+
+
+## Citation
+```
+@inproceedings{wu2018improving,
+  title={Improving Generalization via Scalable Neighborhood Component Analysis},
+  author={Wu, Zhirong and Efros, Alexei A and Yu, Stella},
+  booktitle={European Conference on Computer Vision (ECCV) 2018},
+  year={2018}
+}
+```
+
+## Contact
+
+For any questions, please feel free to reach 
+```
+Zhirong Wu: [email protected]
+```

cifar.py

@@ -0,0 +1,165 @@
+'''Train CIFAR10 with PyTorch.'''
+from __future__ import print_function
+
+import sys
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+
+import torchvision
+import torchvision.transforms as transforms
+
+import os
+import argparse
+import time
+
+import models
+import datasets
+import math
+
+from lib.LinearAverage import LinearAverage
+from lib.NCA import NCACrossEntropy
+from lib.utils import AverageMeter
+from test import NN, kNN
+
+parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
+parser.add_argument('--lr', default=0.1, type=float, help='learning rate')
+parser.add_argument('--resume', '-r', default='', type=str, help='resume from checkpoint')
+parser.add_argument('--test-only', action='store_true', help='test only')
+parser.add_argument('--low-dim', default=128, type=int,
+                    metavar='D', help='feature dimension')
+parser.add_argument('--temperature', default=0.05, type=float,
+                    metavar='T', help='temperature parameter for softmax')
+parser.add_argument('--memory-momentum', default=0.5, type=float,
+                    metavar='M', help='momentum for non-parametric updates')
+
+args = parser.parse_args()
+
+use_cuda = torch.cuda.is_available()
+best_acc = 0  # best test accuracy
+start_epoch = 0  # start from epoch 0 or last checkpoint epoch
+
+# Data
+print('==> Preparing data..')
+transform_train = transforms.Compose([
+    #transforms.RandomCrop(32, padding=4),
+    transforms.RandomResizedCrop(size=32, scale=(0.2,1.)),
+    transforms.RandomGrayscale(p=0.2),
+    transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
+    transforms.RandomHorizontalFlip(),
+    transforms.ToTensor(),
+    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+])
+
+transform_test = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
+])
+
+trainset = datasets.CIFAR10Instance(root='./data', train=True, download=True, transform=transform_train)
+trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
+
+testset = datasets.CIFAR10Instance(root='./data', train=False, download=True, transform=transform_test)
+testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)
+
+classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
+ndata = trainset.__len__()
+
+# Model
+if args.test_only or len(args.resume)>0:
+    # Load checkpoint.
+    print('==> Resuming from checkpoint..')
+    assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
+    checkpoint = torch.load('./checkpoint/'+args.resume)
+    net = checkpoint['net']
+    lemniscate = checkpoint['lemniscate']
+    best_acc = checkpoint['acc']
+    start_epoch = checkpoint['epoch']
+else:
+    print('==> Building model..')
+    net = models.__dict__['ResNet18'](low_dim=args.low_dim)
+    # define leminiscate
+    lemniscate = LinearAverage(args.low_dim, ndata, args.temperature, args.memory_momentum)
+
+# define loss function
+criterion = NCACrossEntropy(torch.LongTensor(trainloader.dataset.train_labels))
+
+if use_cuda:
+    net.cuda()
+    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
+    lemniscate.cuda()
+    criterion.cuda()
+    cudnn.benchmark = True
+
+if args.test_only:
+    acc = kNN(0, net, lemniscate, trainloader, testloader, 30, args.temperature)
+    sys.exit(0)
+
+optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4, nesterov=True)
+
+def adjust_learning_rate(optimizer, epoch):
+    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
+    lr = args.lr * (0.1 ** (epoch // 50))
+    print(lr)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+# Training
+def train(epoch):
+    print('\nEpoch: %d' % epoch)
+    adjust_learning_rate(optimizer, epoch)
+    train_loss = AverageMeter()
+    data_time = AverageMeter()
+    batch_time = AverageMeter()
+    correct = 0
+    total = 0
+
+    # switch to train mode
+    net.train()
+
+    end = time.time()
+    for batch_idx, (inputs, targets, indexes) in enumerate(trainloader):
+        data_time.update(time.time() - end)
+        if use_cuda:
+            inputs, targets, indexes = inputs.cuda(), targets.cuda(), indexes.cuda()
+        optimizer.zero_grad()
+
+        features = net(inputs)
+        outputs = lemniscate(features, indexes)
+        loss = criterion(outputs, indexes)
+
+        loss.backward()
+        optimizer.step()
+
+        train_loss.update(loss.item(), inputs.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        print('Epoch: [{}][{}/{}]'
+              'Time: {batch_time.val:.3f} ({batch_time.avg:.3f}) '
+              'Data: {data_time.val:.3f} ({data_time.avg:.3f}) '
+              'Loss: {train_loss.val:.4f} ({train_loss.avg:.4f})'.format(
+              epoch, batch_idx, len(trainloader), batch_time=batch_time, data_time=data_time, train_loss=train_loss))
+
+for epoch in range(start_epoch, start_epoch+200):
+    train(epoch)
+    acc = kNN(epoch, net, lemniscate, trainloader, testloader, 30, args.temperature)
+
+    if acc > best_acc:
+        print('Saving..')
+        state = {
+            'net': net.module if use_cuda else net,
+            'lemniscate': lemniscate,
+            'acc': acc,
+            'epoch': epoch,
+        }
+        if not os.path.isdir('checkpoint'):
+            os.mkdir('checkpoint')
+        torch.save(state, './checkpoint/ckpt.t7')
+        best_acc = acc
+
+    print('best accuracy: {:.2f}'.format(best_acc*100))

datasets/__init__.py

@@ -0,0 +1,5 @@
+from .folder import ImageFolderInstance
+from .cifar import CIFAR10Instance, CIFAR100Instance
+
+__all__ = ('ImageFolderInstance', 'CIFAR10Instance', 'CIFAR100Instance')
+

datasets/cifar.py

@@ -0,0 +1,42 @@
+from __future__ import print_function
+from PIL import Image
+import torchvision.datasets as datasets
+import torch.utils.data as data
+
+class CIFAR10Instance(datasets.CIFAR10):
+    """CIFAR10Instance Dataset.
+    """
+    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]
+
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+        img = Image.fromarray(img)
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target, index
+
+class CIFAR100Instance(CIFAR10Instance):
+    """CIFAR100Instance Dataset.
+
+    This is a subclass of the `CIFAR10Instance` Dataset.
+    """
+    base_folder = 'cifar-100-python'
+    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
+    filename = "cifar-100-python.tar.gz"
+    tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
+    train_list = [
+        ['train', '16019d7e3df5f24257cddd939b257f8d'],
+    ]
+
+    test_list = [
+        ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
+    ]

datasets/folder.py

@@ -0,0 +1,21 @@
+import torchvision.datasets as datasets
+
+class ImageFolderInstance(datasets.ImageFolder):
+    """: Folder datasets which returns the index of the image as well::
+    """
+    def __getitem__(self, index):
+        """
+        Args:
+            index (int): Index
+        Returns:
+            tuple: (image, target) where target is class_index of the target class.
+        """
+        path, target = self.imgs[index]
+        img = self.loader(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.target_transform is not None:
+            target = self.target_transform(target)
+
+        return img, target, index
+

lib/LinearAverage.py

@@ -0,0 +1,58 @@
+import torch
+from torch.autograd import Function
+from torch import nn
+import math
+
+class LinearAverageOp(Function):
+    @staticmethod
+    def forward(self, x, y, memory, params):
+        T = params[0].item()
+        batchSize = x.size(0)
+
+        # inner product
+        out = torch.mm(x.data, memory.t())
+        out.div_(T) # batchSize * N
+        
+        self.save_for_backward(x, memory, y, params)
+
+        return out
+
+    @staticmethod
+    def backward(self, gradOutput):
+        x, memory, y, params = self.saved_tensors
+        batchSize = gradOutput.size(0)
+        T = params[0].item()
+        momentum = params[1].item()
+        
+        # add temperature
+        gradOutput.data.div_(T)
+
+        # gradient of linear
+        gradInput = torch.mm(gradOutput.data, memory)
+        gradInput.resize_as_(x)
+
+        # update the non-parametric data
+        weight_pos = memory.index_select(0, y.data.view(-1)).resize_as_(x)
+        weight_pos.mul_(momentum)
+        weight_pos.add_(torch.mul(x.data, 1-momentum))
+        w_norm = weight_pos.pow(2).sum(1, keepdim=True).pow(0.5)
+        updated_weight = weight_pos.div(w_norm)
+        memory.index_copy_(0, y, updated_weight)
+        
+        return gradInput, None, None, None
+
+class LinearAverage(nn.Module):
+
+    def __init__(self, inputSize, outputSize, T=0.05, momentum=0.5):
+        super(LinearAverage, self).__init__()
+        stdv = 1 / math.sqrt(inputSize)
+        self.nLem = outputSize
+
+        self.register_buffer('params',torch.tensor([T, momentum]));
+        stdv = 1. / math.sqrt(inputSize/3)
+        self.register_buffer('memory', torch.rand(outputSize, inputSize).mul_(2*stdv).add_(-stdv))
+
+    def forward(self, x, y):
+        out = LinearAverageOp.apply(x, y, self.memory, self.params)
+        return out
+