def __init__(self, pretrained=True, finetune=True):
super(VGGNetFeats, self).__init__()
model = vgg16(pretrained=pretrained)
for param in model.parameters():
param.requires_grad = finetune
self.features = model.features
self.classifier = nn.Sequential(
*list(model.classifier.children())[:-1])
def __init__(self,
image_dir,
sketch_dir,
save_dir,
image_model_path,
sketch_model_path,
tu_berlin=False):
if not tu_berlin:
self.classes = list(TRAIN_CLASS) + list(TEST_CLASS)
else:
self.classes = list(TRAIN_CLASS_TUBERLIN_I) + list(
TEST_CLASS_TUBERLIN_I)
self.image_dir = image_dir
self.sketch_dir = sketch_dir
self.image_model_path = image_model_path
self.sketch_model_path = sketch_model_path
self.save_dir = save_dir
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
self.transform = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor()])
# load sketch model
# self.sketch_model = VGGNetFeats(pretrained=False, finetune=False)
self.sketch_model = vgg16(pretrained=True, return_type=3, dropout=0)
self.load_weight(self.sketch_model, self.sketch_model_path, 'sketch')
for param in self.sketch_model.parameters():
param.requires_grad = False
self.sketch_model = self.sketch_model.cuda()
self.sketch_model.eval()
# load image model
# self.image_model = VGGNetFeats(pretrained=False, finetune=False)
self.image_model = vgg16(pretrained=True, return_type=3, dropout=0)
print(self.image_model)
self.load_weight(self.image_model, self.image_model_path, 'image')
for param in self.image_model.parameters():
param.requires_grad = False
self.image_model = self.image_model.cuda()
self.image_model.eval()
def __init__(self,
margin,
loss_type,
distance_type,
logger,
from_pretrain=True,
batch_normalization=True):
super(Siamese, self).__init__()
self.margin = margin
self.loss_type = loss_type
self.distance_type = distance_type
self.logger = logger
if batch_normalization:
self.model = vgg16_bn(pretrained=from_pretrain, return_type=1)
else:
self.model = vgg16(pretrained=from_pretrain, return_type=1)
def __init__(self, image_dir, sketch_dir, save_dir, dataset='sketchy'):
if dataset == 'sketchy':
self.classes = list(TRAIN_CLASS) + list(TEST_CLASS)
else:
self.classes = list(TRAIN_CLASS_TUBERLIN_I) + list(
TEST_CLASS_TUBERLIN_I)
self.image_dir = image_dir
self.sketch_dir = sketch_dir
self.save_dir = save_dir
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
self.backbone = vgg16(pretrained=True, return_type=3, dropout=0)
for param in self.backbone.parameters():
param.requires_grad = False
self.backbone = self.backbone.cuda()
self.backbone.eval()
self.ToTensor = ToTensor()
self.Normalize = Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
def __init__(self, hidden_size, hashing_bit, semantics_size, pretrain_embedding, adj_scaler=1, dropout=0.5, from_pretrain=True, fix_cnn=True, fix_embedding=True, logger=None):
super(ZSIM, self).__init__()
# hyper-param
self.hidden_size = hidden_size
self.hashing_bit = hashing_bit
self.adj_scaler = adj_scaler
self.dropout = dropout
self.from_pretrain = from_pretrain
self.fix_cnn = fix_cnn
self.logger = logger
# model
self.semantics = nn.Embedding.from_pretrained(pretrain_embedding)
if fix_embedding:
self.semantics.weight.requires_grad=False
self.backbone = vgg16(pretrained=from_pretrain, return_type=1)
if fix_cnn:
for param in self.backbone.parameters():
param.requires_grad = False
self.sketch_attention = CNN_attention(512)
self.sketch_linear1 = nn.Linear(512, hidden_size)
self.sketch_linear2 = nn.Linear(hidden_size, hashing_bit)
self.image_attention = CNN_attention(512)
self.image_linear1 = nn.Linear(512, hidden_size)
self.image_linear2 = nn.Linear(hidden_size, hashing_bit)
self.gcn = GCN_ZSIH(512*512, hidden_size, hashing_bit, dropout, adj_scaler)
self.doubly_sn = Doubly_SN_Function.apply
self.mean_linear = nn.Linear(hashing_bit, semantics_size)
self.var_linear = nn.Linear(hashing_bit, semantics_size)
# activation function
self.relu = nn.ReLU(inplace=True)
self.softmax = nn.Softmax(dim=-1)
self.sigmoid = nn.Sigmoid()
# loss function
self.l2 = L2()
self.bce = CE()