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)
gitextract_vz9ijhsx/ ├── .gitignore ├── LICENSE ├── README.md ├── evaluate.py ├── net.py └── train.py
SYMBOL INDEX (9 symbols across 3 files)
FILE: evaluate.py
function load_data (line 21) | def load_data():
function binary_output (line 42) | def binary_output(dataloader):
function evaluate (line 59) | def evaluate(trn_binary, trn_label, tst_binary, tst_label):
FILE: net.py
class AlexNetPlusLatent (line 9) | class AlexNetPlusLatent(nn.Module):
method __init__ (line 10) | def __init__(self, bits):
method forward (line 19) | def forward(self, x):
FILE: train.py
function init_dataset (line 30) | def init_dataset():
function train (line 53) | def train(epoch_num):
function test (line 77) | def test():
Condensed preview — 6 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (14K chars).
[
{
"path": ".gitignore",
"chars": 7,
"preview": "/data/\n"
},
{
"path": "LICENSE",
"chars": 1067,
"preview": "MIT License\n\nCopyright (c) 2020 Fan Jingbo\n\nPermission is hereby granted, free of charge, to any person obtaining a copy"
},
{
"path": "README.md",
"chars": 709,
"preview": "# pytorch_deephash\n\n## Introduction\n\nThis is the Pytorch implementation of [Deep Learning of Binary Hash Codes for Fast "
},
{
"path": "evaluate.py",
"chars": 5732,
"preview": "import argparse\nimport os\nfrom timeit import time\n\nimport numpy as np\nimport torch\nimport torch.optim.lr_scheduler\nfrom "
},
{
"path": "net.py",
"chars": 851,
"preview": "import os\nimport torch.nn as nn\nfrom torchvision import models\n\nos.environ['TORCH_HOME'] = 'models'\nalexnet_model = mode"
},
{
"path": "train.py",
"chars": 5138,
"preview": "import argparse\nimport math\nimport os\nimport shutil\n\nimport torch\nimport torch.nn as nn\nimport torch.optim.lr_scheduler\n"
}
]
About this extraction
This page contains the full source code of the flyingpot/pytorch_deephash GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 6 files (13.2 KB), approximately 3.3k tokens, and a symbol index with 9 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.