From 764e950f295ee25a02f12c13a8379b778f24a4f4 Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 11:54:25 +0900 Subject: [PATCH 1/8] Added a supervised metric pretraining option as described in the paper --- lib/SupervisedMetricPretraining.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100755 lib/SupervisedMetricPretraining.py diff --git a/lib/SupervisedMetricPretraining.py b/lib/SupervisedMetricPretraining.py new file mode 100755 index 0000000..15670ca --- /dev/null +++ b/lib/SupervisedMetricPretraining.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Wed Jan 27 11:06:53 2021 + +@author: nuvilabs +""" + +class NuviS3KeySensor(S3KeySensor): + + + def poke(self, context): + to_return = super().poke(context) + + if to_return: + hook = S3Hook(aws_conn_id=self.aws_conn_id) + keys = hook.list_keys(self.bucket_name) + + check_key = lambda key: 'ID' in key + check_key_v = np.vectorize(check_key) + keys = np.asarray(keys) + idx = check_key_v(keys) + ID_keys = keys[idx].tolist() + to_return = len(ID_keys) != 0 + if to_return:self.xcom_push(context, key="IDS", value = ID_keys) + + return to_return \ No newline at end of file From 67fe4e799b3ce43a09b5ae641f048f6748b59c89 Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 11:57:51 +0900 Subject: [PATCH 2/8] Added a supervised metric pretraining option as described in the paper --- lib/SupervisedMetricPretraining.py | 85 ++++++++++++++++++++++-------- 1 file changed, 64 insertions(+), 21 deletions(-) diff --git a/lib/SupervisedMetricPretraining.py b/lib/SupervisedMetricPretraining.py index 15670ca..d12fcd3 100755 --- a/lib/SupervisedMetricPretraining.py +++ b/lib/SupervisedMetricPretraining.py @@ -1,27 +1,70 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -Created on Wed Jan 27 11:06:53 2021 +import torch +from torch import nn +train_labels_= None +def get_train_labels(trainloader, device='cuda'): + global train_labels_ + if train_labels_ is None: + print("=> loading all train labels") + train_labels = -1 * torch.ones([len(trainloader.dataset)], dtype=torch.long) + for i, (_, label, index) in enumerate(trainloader): + train_labels[index] = label + if i % 10000 == 0: + print("{}/{}".format(i, len(trainloader))) + assert all(train_labels != -1) + train_labels_ = train_labels.to(device) + return train_labels_ +class Supervised_Pretraining(object): + def __init__(self,trainloader,n, t=0.07): + """ + Parameters + ---------- + trainloader : + DataLoader containing training data. + n : int + Number of labels. + t : float + Temperature parameter as described in https://arxiv.org/pdf/1805.01978.pdf. -@author: nuvilabs -""" + """ + super(Supervised_Pretraining,self).__init__() + # get train labels + self.labels = get_train_labels(trainloader) + # Softmax loss + self.loss_fn = nn.CrossEntropyLoss() + #init labels + self.n_labels = n + self.t = t + def to(self,device): + #send to a device + self.loss_fn.to(device) + def __call__(self,out,y): + return self.forward(out,y) + def forward(self,out,y): + """ + Parameters + ---------- + out : + Output from LinearAverage.py as described in https://arxiv.org/pdf/1812.08781.pdf. + y : tensor + Target Labels. -class NuviS3KeySensor(S3KeySensor): + + Returns + ------- + Softmax Loss. - - def poke(self, context): - to_return = super().poke(context) + """ + #making it more sensitive by dividing by temperature value as in https://arxiv.org/pdf/1805.01978.pdf + out.div_(self.t) + #eq (4) in https://arxiv.org/pdf/1812.08781.pdf + scores = torch.zeros(out.shape[0],self.n_labels).cuda() + for i in range(self.n_labels): + yi = self.labels == i - if to_return: - hook = S3Hook(aws_conn_id=self.aws_conn_id) - keys = hook.list_keys(self.bucket_name) + candidates = yi.view(1,-1).expand(out.shape[0], -1) + retrieval = out[candidates] + retrieval = retrieval.reshape(out.shape[0], -1) - check_key = lambda key: 'ID' in key - check_key_v = np.vectorize(check_key) - keys = np.asarray(keys) - idx = check_key_v(keys) - ID_keys = keys[idx].tolist() - to_return = len(ID_keys) != 0 - if to_return:self.xcom_push(context, key="IDS", value = ID_keys) + scores[:,i] = retrieval.sum(1,keepdim=True).view(1,-1) - return to_return \ No newline at end of file + return self.loss_fn(scores, y) \ No newline at end of file From 7a10c2964a3c0b66a3344d204df8bc8cd4a6d0db Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 22:03:25 +0900 Subject: [PATCH 3/8] supervised metric pretraining --- lib/SupervisedMetricPretraining.py | 70 ++------ supervised/imagenet.py | 268 +++++++++++++++++++++++++++++ 2 files changed, 285 insertions(+), 53 deletions(-) create mode 100755 supervised/imagenet.py diff --git a/lib/SupervisedMetricPretraining.py b/lib/SupervisedMetricPretraining.py index d12fcd3..afd02a5 100755 --- a/lib/SupervisedMetricPretraining.py +++ b/lib/SupervisedMetricPretraining.py @@ -1,70 +1,34 @@ import torch from torch import nn -train_labels_= None -def get_train_labels(trainloader, device='cuda'): - global train_labels_ - if train_labels_ is None: - print("=> loading all train labels") - train_labels = -1 * torch.ones([len(trainloader.dataset)], dtype=torch.long) - for i, (_, label, index) in enumerate(trainloader): - train_labels[index] = label - if i % 10000 == 0: - print("{}/{}".format(i, len(trainloader))) - assert all(train_labels != -1) - train_labels_ = train_labels.to(device) - return train_labels_ -class Supervised_Pretraining(object): - def __init__(self,trainloader,n, t=0.07): - """ - Parameters - ---------- - trainloader : - DataLoader containing training data. - n : int - Number of labels. - t : float - Temperature parameter as described in https://arxiv.org/pdf/1805.01978.pdf. +from .utils import get_train_labels - """ - super(Supervised_Pretraining,self).__init__() +class SupervisedSoftmax(object): + def __init__(self,trainloader,device,t=0.07): + super(SupervisedSoftmax,self).__init__() # get train labels self.labels = get_train_labels(trainloader) # Softmax loss - self.loss_fn = nn.CrossEntropyLoss() + self.loss_fn = nn.CrossEntropyLoss().to(device) #init labels - self.n_labels = n - self.t = t + self.n_labels = self.labels.max().data.item() + 1 + print(self.n_labels) + #Temperature parameter as described in https://arxiv.org/pdf/1805.01978.pdf. + self.temperature = t def to(self,device): #send to a device self.loss_fn.to(device) def __call__(self,out,y): return self.forward(out,y) - def forward(self,out,y): - """ - Parameters - ---------- - out : - Output from LinearAverage.py as described in https://arxiv.org/pdf/1812.08781.pdf. - y : tensor - Target Labels. - - - Returns - ------- - Softmax Loss. - - """ + def forward(self,dist,y): #making it more sensitive by dividing by temperature value as in https://arxiv.org/pdf/1805.01978.pdf - out.div_(self.t) + dist.div_(self.temperature) #eq (4) in https://arxiv.org/pdf/1812.08781.pdf - scores = torch.zeros(out.shape[0],self.n_labels).cuda() + scores = torch.zeros(dist.shape[0],self.n_labels).cuda() for i in range(self.n_labels): yi = self.labels == i - - candidates = yi.view(1,-1).expand(out.shape[0], -1) - retrieval = out[candidates] - retrieval = retrieval.reshape(out.shape[0], -1) - + candidates = yi.view(1,-1).expand(dist.shape[0], -1) + retrieval = dist[candidates] + retrieval = retrieval.reshape(dist.shape[0], -1) scores[:,i] = retrieval.sum(1,keepdim=True).view(1,-1) - - return self.loss_fn(scores, y) \ No newline at end of file + + return self.loss_fn(scores, y) diff --git a/supervised/imagenet.py b/supervised/imagenet.py new file mode 100755 index 0000000..f7ee3d7 --- /dev/null +++ b/supervised/imagenet.py @@ -0,0 +1,268 @@ +import argparse +import os +import shutil +import time + +import torch +import torch.nn as nn +import torch.nn.parallel +import torch.backends.cudnn as cudnn +import torch.optim +import torch.utils.data +import torchvision.transforms as transforms + +from lib import models, datasets + +from lib.NCEAverage import NCEAverage +from lib.LinearAverage import LinearAverage +from lib.NCECriterion import NCECriterion +from lib.utils import AverageMeter +from lib.SupervisedMetricPretraining import SupervisedSoftmax +from test import NN, kNN + +model_names = sorted(name for name in models.__dict__ + if name.islower() and not name.startswith("__") + and callable(models.__dict__[name])) + +parser = argparse.ArgumentParser(description='PyTorch ImageNet Training') +parser.add_argument('--data-dir', metavar='DIR', + help='path to dataset', required=True) +parser.add_argument('--model-dir', metavar='DIR', + default='./checkpoint/instance_imagenet', help='path to save model') +parser.add_argument('--log-dir', metavar='DIR', + default='./tensorboard/instance_imagenet', help='path to save log') +parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18', + choices=model_names, + help='model architecture: ' + + ' | '.join(model_names) + + ' (default: resnet18)') +parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', + help='number of data loading workers (default: 4)') +parser.add_argument('--epochs', default=200, type=int, metavar='N', + help='number of total epochs to run') +parser.add_argument('--start-epoch', default=0, type=int, metavar='N', + help='manual epoch number (useful on restarts)') +parser.add_argument('-b', '--batch-size', default=256, type=int, + metavar='N', help='mini-batch size (default: 256)') +parser.add_argument('-vb', '--val-batch-size', default=128, type=int, + metavar='N', help='validation mini-batch size (default: 128)') +parser.add_argument('--lr', '--learning-rate', default=0.03, type=float, + metavar='LR', help='initial learning rate') +parser.add_argument('--momentum', default=0.9, type=float, metavar='M', + help='momentum') +parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, + metavar='W', help='weight decay (default: 1e-4)') +parser.add_argument('--print-freq', '-p', default=10, type=int, + metavar='N', help='print frequency (default: 10)') +parser.add_argument('--resume', default='', type=str, metavar='PATH', + help='path to latest checkpoint (default: none)') +parser.add_argument('--auto-resume', action='store_true', help='auto resume') +parser.add_argument('--test-only', action='store_true', help='test only') +parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', + help='evaluate model on validation set') +parser.add_argument('--pretrained', dest='pretrained', action='store_true', + help='use pre-trained model') +parser.add_argument('--low-dim', default=128, type=int, + metavar='D', help='feature dimension') +parser.add_argument('--nce-k', default=4096, type=int, + metavar='K', help='number of negative samples for NCE') +parser.add_argument('--nce-t', default=0.07, type=float, + metavar='T', help='temperature parameter for softmax') +parser.add_argument('--nce-m', default=0.5, type=float, + help='momentum for non-parametric updates') +parser.add_argument('--iter-size', default=1, type=int, + help='caffe style iter size') + +best_prec1 = 0 + + +def main(): + global args, best_prec1 + args = parser.parse_args() + + # create model + if args.pretrained: + print("=> using pre-trained model '{}'".format(args.arch)) + model = models.__dict__[args.arch](pretrained=True) + else: + print("=> creating model '{}'".format(args.arch)) + model = models.__dict__[args.arch](low_dim=args.low_dim) + + if args.arch.startswith('alexnet') or args.arch.startswith('vgg'): + model.features = torch.nn.DataParallel(model.features) + model.cuda() + else: + model = torch.nn.DataParallel(model).cuda() + + # Data loading code + print("=> loading dataset") + traindir = os.path.join(args.data_dir, 'train') + valdir = os.path.join(args.data_dir, 'val') + normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + + train_dataset = datasets.ImageFolderInstance( + traindir, + transforms.Compose([ + transforms.RandomResizedCrop(224, scale=(0.2, 1.)), + transforms.RandomGrayscale(p=0.2), + transforms.ColorJitter(0.4, 0.4, 0.4, 0.4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + normalize, + ])) + + train_loader = torch.utils.data.DataLoader( + train_dataset, batch_size=args.batch_size, shuffle=True, + num_workers=args.workers, pin_memory=True) + + val_loader = torch.utils.data.DataLoader( + datasets.ImageFolderInstance(valdir, transforms.Compose([ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + normalize, + ])), + batch_size=args.val_batch_size, shuffle=False, + num_workers=args.workers, pin_memory=True) + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + # define lemniscate and loss function (criterion) + print("=> building optimizer") + ndata = train_dataset.__len__() + + lemniscate = LinearAverage( + args.low_dim, ndata, args.nce_t, args.nce_m).cuda() + criterion = SupervisedSoftmax(train_loader,device) + + optimizer = torch.optim.SGD(model.parameters(), args.lr, + momentum=args.momentum, + weight_decay=args.weight_decay) + + # optionally resume from a checkpoint + model_filename_to_resume = None + if args.resume: + if os.path.isfile(args.resume): + model_filename_to_resume = args.resume + else: + print("=> no checkpoint found at '{}'".format(args.resume)) + elif args.auto_resume: + for epoch in range(args.epochs, args.start_epoch + 1, -1): + model_filename = get_model_name(epoch) + if os.path.exists(model_filename): + model_filename_to_resume = model_filename + break + else: + print("=> no checkpoint found at '{}'".format(args.model_dir)) + + if model_filename_to_resume is not None: + print("=> loading checkpoint '{}'".format(model_filename_to_resume)) + checkpoint = torch.load(model_filename_to_resume) + args.start_epoch = checkpoint['epoch'] + best_prec1 = checkpoint['best_prec1'] + model.load_state_dict(checkpoint['state_dict']) + lemniscate = checkpoint['lemniscate'] + optimizer.load_state_dict(checkpoint['optimizer']) + print("=> loaded checkpoint '{}' (epoch {})" + .format(model_filename_to_resume, checkpoint['epoch'])) + + cudnn.benchmark = True + + if args.evaluate: + kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t) + return + + for epoch in range(args.start_epoch, args.epochs): + adjust_learning_rate(optimizer, epoch) + + # train for one epoch + train(train_loader, model, lemniscate, criterion, optimizer, epoch) + + # evaluate on validation set + prec1 = NN(model, lemniscate, train_loader, val_loader) + + # remember best prec@1 and save checkpoint + is_best = prec1 > best_prec1 + best_prec1 = max(prec1, best_prec1) + + save_checkpoint({ + 'epoch': epoch + 1, + 'arch': args.arch, + 'state_dict': model.state_dict(), + 'lemniscate': lemniscate, + 'best_prec1': best_prec1, + 'optimizer': optimizer.state_dict(), + }, is_best, + filename=get_model_name(epoch)) + + # evaluate KNN after last epoch + kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t) + + +def train(train_loader, model, lemniscate, criterion, optimizer, epoch): + batch_time = AverageMeter() + data_time = AverageMeter() + losses = AverageMeter() + + # switch to train mode + model.train() + + end = time.time() + optimizer.zero_grad() + for i, (inputs, target, index) in enumerate(train_loader): + # measure data loading time + data_time.update(time.time() - end) + + index = index.cuda(non_blocking=True) + target = target.cuda(non_blocking=True) + # compute output + feature = model(inputs) + output = lemniscate(feature, index) + loss = criterion(output, target) + + loss.backward() + + # measure accuracy and record loss + losses.update(loss.item() * args.iter_size, inputs.size(0)) + + if (i + 1) % args.iter_size == 0: + # compute gradient and do SGD step + optimizer.step() + optimizer.zero_grad() + + # measure elapsed time + batch_time.update(time.time() - end) + end = time.time() + + if i % args.print_freq == 0: + print(f'Epoch: [{epoch}/{args.epochs}][{i}/{len(train_loader)}]\t' + f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' + f'Data {data_time.val:.3f} ({data_time.avg:.3f})\t' + f'Loss {losses.val:.4f} ({losses.avg:.4f})\t') + + +def get_model_name(epoch): + return os.path.join(args.model_dir, 'ckpt-{}.pth.tar'.format(epoch)) + + +def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'): + torch.save(state, filename) + if is_best: + shutil.copyfile(filename, os.path.join( + args.model_dir, 'model_best.pth.tar')) + + +def adjust_learning_rate(optimizer, epoch): + """Sets the learning rate to the initial LR decayed by 10 every 100 epochs""" + if epoch < 120: + lr = args.lr + elif 120 <= epoch < 160: + lr = args.lr * 0.1 + else: + lr = args.lr * 0.01 + # lr = args_.lr * (0.1 ** (epoch // 100)) + for param_group in optimizer.param_groups: + param_group['lr'] = lr + + +if __name__ == '__main__': + main() From eb9c4fd4fbc3321f15314376fc3e817637ce2824 Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 22:07:59 +0900 Subject: [PATCH 4/8] fix: out -> dist --- lib/SupervisedMetricPretraining.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/SupervisedMetricPretraining.py b/lib/SupervisedMetricPretraining.py index afd02a5..afce608 100755 --- a/lib/SupervisedMetricPretraining.py +++ b/lib/SupervisedMetricPretraining.py @@ -17,8 +17,8 @@ def __init__(self,trainloader,device,t=0.07): def to(self,device): #send to a device self.loss_fn.to(device) - def __call__(self,out,y): - return self.forward(out,y) + def __call__(self,dist,y): + return self.forward(dist,y) def forward(self,dist,y): #making it more sensitive by dividing by temperature value as in https://arxiv.org/pdf/1805.01978.pdf dist.div_(self.temperature) From ebcd8417d0ddd93a87b68960d2f8776d0b378b33 Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 22:19:27 +0900 Subject: [PATCH 5/8] Update SupervisedMetricPretraining.py --- lib/SupervisedMetricPretraining.py | 1 - 1 file changed, 1 deletion(-) diff --git a/lib/SupervisedMetricPretraining.py b/lib/SupervisedMetricPretraining.py index afce608..71cefbd 100755 --- a/lib/SupervisedMetricPretraining.py +++ b/lib/SupervisedMetricPretraining.py @@ -11,7 +11,6 @@ def __init__(self,trainloader,device,t=0.07): self.loss_fn = nn.CrossEntropyLoss().to(device) #init labels self.n_labels = self.labels.max().data.item() + 1 - print(self.n_labels) #Temperature parameter as described in https://arxiv.org/pdf/1805.01978.pdf. self.temperature = t def to(self,device): From 3efa84ce124e96c08eb0de5b2d4ec58706f07468 Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 22:22:41 +0900 Subject: [PATCH 6/8] fix: remove nce_k for supervised --- supervised/imagenet.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/supervised/imagenet.py b/supervised/imagenet.py index f7ee3d7..5997062 100755 --- a/supervised/imagenet.py +++ b/supervised/imagenet.py @@ -64,8 +64,7 @@ help='use pre-trained model') parser.add_argument('--low-dim', default=128, type=int, metavar='D', help='feature dimension') -parser.add_argument('--nce-k', default=4096, type=int, - metavar='K', help='number of negative samples for NCE') + parser.add_argument('--nce-t', default=0.07, type=float, metavar='T', help='temperature parameter for softmax') parser.add_argument('--nce-m', default=0.5, type=float, From f05c625cee938fb02ee02590878ac4c8af2c7a73 Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 22:22:58 +0900 Subject: [PATCH 7/8] Update imagenet.py --- supervised/imagenet.py | 1 - 1 file changed, 1 deletion(-) diff --git a/supervised/imagenet.py b/supervised/imagenet.py index 5997062..53881f0 100755 --- a/supervised/imagenet.py +++ b/supervised/imagenet.py @@ -64,7 +64,6 @@ help='use pre-trained model') parser.add_argument('--low-dim', default=128, type=int, metavar='D', help='feature dimension') - parser.add_argument('--nce-t', default=0.07, type=float, metavar='T', help='temperature parameter for softmax') parser.add_argument('--nce-m', default=0.5, type=float, From 4fa01543d4156d03560ab4e8423332708829ec93 Mon Sep 17 00:00:00 2001 From: Chingis Oinar Date: Wed, 27 Jan 2021 22:42:58 +0900 Subject: [PATCH 8/8] fix:format --- lib/SupervisedMetricPretraining.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/lib/SupervisedMetricPretraining.py b/lib/SupervisedMetricPretraining.py index 71cefbd..9d93147 100755 --- a/lib/SupervisedMetricPretraining.py +++ b/lib/SupervisedMetricPretraining.py @@ -3,6 +3,7 @@ from .utils import get_train_labels class SupervisedSoftmax(object): + """Supervised Metric Pretraining.""" def __init__(self,trainloader,device,t=0.07): super(SupervisedSoftmax,self).__init__() # get train labels @@ -13,11 +14,14 @@ def __init__(self,trainloader,device,t=0.07): self.n_labels = self.labels.max().data.item() + 1 #Temperature parameter as described in https://arxiv.org/pdf/1805.01978.pdf. self.temperature = t + def to(self,device): #send to a device self.loss_fn.to(device) + def __call__(self,dist,y): return self.forward(dist,y) + def forward(self,dist,y): #making it more sensitive by dividing by temperature value as in https://arxiv.org/pdf/1805.01978.pdf dist.div_(self.temperature)