Repository: flyingpot/pytorch_deephash Branch: master Commit: fbd94a2cfe1c Files: 6 Total size: 13.2 KB Directory structure: gitextract_vz9ijhsx/ ├── .gitignore ├── LICENSE ├── README.md ├── evaluate.py ├── net.py └── train.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ /data/ ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2020 Fan Jingbo Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: README.md ================================================ # pytorch_deephash ## Introduction This is the Pytorch implementation of [Deep Learning of Binary Hash Codes for Fast Image Retrieval](https://github.com/kevinlin311tw/caffe-cvprw15), and can achieve more than 93% mAP in CIFAR10 dataset. ## Environment > Pytorch 1.4.0 > > torchvision 0.5.0 > > tqdm > > numpy ## Training ```bash python train.py ``` You will get trained models in model folder by default, and models' names are their test accuracy. ## Evaluation ```bash python evaluate.py --pretrained {your saved model name in model folder by default} ``` ## Tips 1. If using Windows, keep num_works zero 2. There are some other args, which you can get them by adding '-h' or reading the code. ================================================ FILE: evaluate.py ================================================ import argparse import os from timeit import time import numpy as np import torch import torch.optim.lr_scheduler from torchvision import datasets, transforms from tqdm import tqdm from net import AlexNetPlusLatent parser = argparse.ArgumentParser(description='Deep Hashing evaluate mAP') parser.add_argument('--pretrained', type=float, default=0, metavar='pretrained_model', help='loading pretrained model(default = None)') parser.add_argument('--bits', type=int, default=48, metavar='bts', help='binary bits') args = parser.parse_args() def load_data(): transform_train = transforms.Compose( [transforms.Resize(227), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))]) transform_test = transforms.Compose( [transforms.Resize(227), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))]) trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=100, shuffle=False, num_workers=0) testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=0) return trainloader, testloader def binary_output(dataloader): net = AlexNetPlusLatent(args.bits) net.load_state_dict(torch.load('./model/{}'.format(args.pretrained))) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print("Use device: " + str(device)) net.to(device) full_batch_output = torch.cuda.FloatTensor() full_batch_label = torch.cuda.LongTensor() with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(dataloader): inputs, targets = inputs.to(device), targets.to(device) outputs, _ = net(inputs) full_batch_output = torch.cat((full_batch_output, outputs.data), 0) full_batch_label = torch.cat((full_batch_label, targets.data), 0) return torch.round(full_batch_output), full_batch_label def evaluate(trn_binary, trn_label, tst_binary, tst_label): classes = np.max(tst_label) + 1 for i in range(classes): if i == 0: tst_sample_binary = tst_binary[np.random.RandomState(seed=i).permutation(np.where(tst_label == i)[0])[:100]] tst_sample_label = np.array([i]).repeat(100) continue else: tst_sample_binary = np.concatenate([tst_sample_binary, tst_binary[np.random.RandomState(seed=i).permutation(np.where(tst_label==i)[0])[:100]]]) tst_sample_label = np.concatenate([tst_sample_label, np.array([i]).repeat(100)]) query_times = tst_sample_binary.shape[0] trainset_len = trn_binary.shape[0] AP = np.zeros(query_times) precision_radius = np.zeros(query_times) Ns = np.arange(1, trainset_len + 1) sum_tp = np.zeros(trainset_len) total_time_start = time.time() with tqdm(total=query_times, desc="Query") as pbar: for i in range(query_times): query_label = tst_sample_label[i] query_binary = tst_sample_binary[i, :] query_result = np.count_nonzero(query_binary != trn_binary, axis=1) # don't need to divide binary length sort_indices = np.argsort(query_result) buffer_yes = np.equal(query_label, trn_label[sort_indices]).astype(int) P = np.cumsum(buffer_yes) / Ns precision_radius[i] = P[np.where(np.sort(query_result) > 2)[0][0]-1] AP[i] = np.sum(P * buffer_yes) / sum(buffer_yes) sum_tp = sum_tp + np.cumsum(buffer_yes) pbar.set_postfix({'Average Precision': '{0:1.5f}'.format(AP[i])}) pbar.update(1) pbar.close() mAP = np.mean(AP) precision_at_k = sum_tp / Ns / query_times index = [100, 200, 400, 600, 800, 1000] index = [i - 1 for i in index] print('precision at k:', precision_at_k[index]) print('precision within Hamming radius 2:', np.mean(precision_radius)) map = np.mean(AP) print('mAP:', map) print('Total query time:', time.time() - total_time_start) if __name__ == "__main__": if os.path.exists('./result/train_binary') and os.path.exists('./result/train_label') and \ os.path.exists('./result/test_binary') and os.path.exists('./result/test_label') and args.pretrained == 0: train_binary = torch.load('./result/train_binary') train_label = torch.load('./result/train_label') test_binary = torch.load('./result/test_binary') test_label = torch.load('./result/test_label') else: trainloader, testloader = load_data() train_binary, train_label = binary_output(trainloader) test_binary, test_label = binary_output(testloader) if not os.path.isdir('result'): os.mkdir('result') torch.save(train_binary, './result/train_binary') torch.save(train_label, './result/train_label') torch.save(test_binary, './result/test_binary') torch.save(test_label, './result/test_label') train_binary = train_binary.cpu().numpy() train_binary = np.asarray(train_binary, np.int32) train_label = train_label.cpu().numpy() test_binary = test_binary.cpu().numpy() test_binary = np.asarray(test_binary, np.int32) test_label = test_label.cpu().numpy() evaluate(train_binary, train_label, test_binary, test_label) ================================================ FILE: net.py ================================================ import os import torch.nn as nn from torchvision import models os.environ['TORCH_HOME'] = 'models' alexnet_model = models.alexnet(pretrained=True) class AlexNetPlusLatent(nn.Module): def __init__(self, bits): super(AlexNetPlusLatent, self).__init__() self.bits = bits self.features = nn.Sequential(*list(alexnet_model.features.children())) self.remain = nn.Sequential(*list(alexnet_model.classifier.children())[:-1]) self.Linear1 = nn.Linear(4096, self.bits) self.sigmoid = nn.Sigmoid() self.Linear2 = nn.Linear(self.bits, 10) def forward(self, x): x = self.features(x) x = x.view(x.size(0), 256 * 6 * 6) x = self.remain(x) x = self.Linear1(x) features = self.sigmoid(x) result = self.Linear2(features) return features, result ================================================ FILE: train.py ================================================ import argparse import math import os import shutil import torch import torch.nn as nn import torch.optim.lr_scheduler from torchvision import datasets, transforms from tqdm import tqdm from net import AlexNetPlusLatent parser = argparse.ArgumentParser(description='Deep Hashing') parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='SGD momentum (default: 0.9)') parser.add_argument('--epoch', type=int, default=128, metavar='epoch', help='epoch') parser.add_argument('--pretrained', type=str, default=0, metavar='pretrained_model', help='loading pretrained model(default = None)') parser.add_argument('--bits', type=int, default=48, metavar='bts', help='binary bits') parser.add_argument('--path', type=str, default='model', metavar='P', help='path directory') args = parser.parse_args() def init_dataset(): transform_train = transforms.Compose( [transforms.Resize(256), transforms.RandomCrop(227), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))]) transform_test = transforms.Compose( [transforms.Resize(227), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))]) trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=0) testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=True, num_workers=0) return trainloader, testloader def train(epoch_num): print('\nEpoch: %d' % epoch_num) net.train() train_loss = 0 correct = 0 total = 0 with tqdm(total=math.ceil(len(trainloader)), desc="Training") as pbar: for batch_idx, (inputs, targets) in enumerate(trainloader): inputs, targets = inputs.to(device), targets.to(device) _, outputs = net(inputs) loss = softmaxloss(outputs, targets) optimizer4nn.zero_grad() loss.backward() optimizer4nn.step() train_loss += softmaxloss(outputs, targets).item() _, predicted = torch.max(outputs.data, 1) total += targets.size(0) correct += predicted.eq(targets.data).sum() pbar.set_postfix({'loss': '{0:1.5f}'.format(loss), 'accurate': '{:.2%}'.format(correct.item() / total)}) pbar.update(1) pbar.close() return train_loss / (batch_idx + 1) def test(): net.eval() with torch.no_grad(): test_loss = 0 correct = 0 total = 0 with tqdm(total=math.ceil(len(testloader)), desc="Testing") as pbar: for batch_idx, (inputs, targets) in enumerate(testloader): inputs, targets = inputs.to(device), targets.to(device) _, outputs = net(inputs) loss = softmaxloss(outputs, targets) test_loss += loss.item() _, predicted = torch.max(outputs.data, 1) total += targets.size(0) correct += predicted.eq(targets.data).sum() pbar.set_postfix({'loss': '{0:1.5f}'.format(loss), 'accurate': '{:.2%}'.format(correct.item() / total)}) pbar.update(1) pbar.close() acc = 100 * int(correct) / int(total) if epoch == args.epoch: print('Saving') if not os.path.isdir('{}'.format(args.path)): os.mkdir('{}'.format(args.path)) torch.save(net.state_dict(), './{}/{}'.format(args.path, acc)) if __name__ == '__main__': torch.cuda.empty_cache() # When using windows, this line is needed trainloader, testloader = init_dataset() net = AlexNetPlusLatent(args.bits) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print("Use device: " + str(device)) net.to(device) softmaxloss = nn.CrossEntropyLoss().cuda() optimizer4nn = torch.optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=0.0005) scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer4nn, milestones=[args.epoch], gamma=0.1) best_acc = 0 start_epoch = 1 if args.pretrained: net.load_state_dict(torch.load('./{}/{}'.format(args.path, args.pretrained))) test() else: if os.path.isdir('{}'.format(args.path)): shutil.rmtree('{}'.format(args.path)) for epoch in range(start_epoch, start_epoch + args.epoch): train(epoch) test() scheduler.step(epoch)