有很多的检索相关的任务是使用哈希码的方式保存数据,然后在检索的时候将哈希码进行汉明距比较得出相似度的判断,那么如何去完成对于哈希码的保存是非常重要的一个环节。
博主最近学习了如何去完成对于哈希码的保存,也是踩了很多坑。代码还在跑没事干,那么将哈希码生成步骤记录一下,分享给大家参考的同时也作为一个笔记,方便以后回顾。
那么废话不多说就开始,哈希码的生成包括有以下五个步骤:
1. 网络的训练
首先要对原有的主干网络进行训练。只有网络训练后能够很好的提取特征了才可以凝练出有用的哈希码。所以可以将哈希码的生成放在整个运行代码的最后部分来实现。
2. 置空全连接层
删除全连接层的步骤是因为一般训练分类的网络,全连接层后得出的结果是分类的结果。比如,你在分类的时候一共有10个类,那么全连接层一般最后输出的shape就是10。
而哈希码通过人为设置,可以是8,16,32,64等。所以需要将全连接层进行删除,并在之后的哈希网络中使用自己定义的哈希码长度来生成哈希码。
下面是删除全连接层的代码,这个代码我踩了很多坑,后来发现只有置空才是最稳定的,del 这类的直接删除会倒置后面在运行的过程中前向运行出现问题。
# 置空全连接层
# fc2是训练模型的最后一个全连接层的名称
model.fc2 = nn.Identity()
3.设置生成哈希网络
然后是设置生成哈希网络,这个网络既要承接上面训练的网络结构,又要生成需要的哈希码,代码如下所示。
import torch
import torch.nn as nn
DEVICE = torch.device("cuda")
# 哈希码的长度设置为32
# 训练模型的置空层上一层的输出为1000
class HashingNet(nn.Module):
def __init__(self, pretrained_model, hash_bits):
super(HashingNet, self).__init__()
self.pretrained_model = pretrained_model
self.hash_layer = nn.Linear(1000, hash_bits)
def forward(self, x):
features = self.pretrained_model(x)
hash_code = self.hash_layer(features)
hash_code_sgn = torch.sign(hash_code)
return hash_code_sgn
使用全连接层和sign函数就可以实现对于哈希码的构建。
下面是在运行代码中对于哈希网络的调用,创建了一个模型框架为哈希网络的模型。
model_hash = HashingNet(model,hash_bits=32)
4. 转化为哈希码
创建完哈希网络之后就是转化为哈希码,哈希码的生成还需要使用到训练集中的数据,因为哈希的意义就在于存储之前的数据然后对查询的数据进行比较。
注意哈希网络也是模型,需要进行to(DEVICE)操作,之前没有进行这个操作报错了我一直没有找到哪里错误哈哈哈哈哈哈。
下面是转化为哈希码的代码:
model_hash.to(DEVICE)
model_hash.eval()
with torch.no_grad():
for data, target in train_loader:
hash_codes = model_hash(data.to(DEVICE))
hash_codes_np = hash_codes.cpu().numpy()
5. 将生成的哈希码保存到txt文件中
最后将生成的哈希码保存到txt文件中,下面是转化为哈希码的代码:
np.savetxt('hash_codes.txt', hash_codes_np," ",target, fmt='%d')