def forward(self, data_shot, data_query, train=True):
proto = self.encoder(data_shot)
proto_q = self.encoder(data_query)
if train:
proto_glvq = self.glvq(proto, self.encoder)
proto_q_glvq = self.glvq(proto_q, self.encoder)
logits = euclidean_metric(proto_q_glvq,
proto_glvq) / self.args.temperature
else:
proto = proto.reshape(self.args.shot, self.args.way,
-1).mean(dim=0)
logits = euclidean_metric(proto_q, proto) / self.args.temperature
return logits
def forward_proto(self, data_shot, data_query, way = None):
if way is None:
way = self.args.num_class
proto = self.encoder(data_shot)
proto = proto.reshape(self.args.shot, way, -1).mean(dim=0)
query = self.encoder(data_query)
logits = euclidean_metric(query, proto)
return logits
def forward(self, support, query, mode='test'):
# feature extraction
support = self.encoder(support)
# get mean of the support
proto = support.reshape(self.args.shot, -1,
support.shape[-1]).mean(dim=0) # N x d
num_proto = proto.shape[0]
# for query set
query = self.encoder(query)
# adapt the support set instances
proto = proto.unsqueeze(0) # 1 x N x d
# refine by Transformer
proto = self.slf_attn(proto, proto, proto)
proto = proto.squeeze(0)
# compute distance for all batches
logitis = euclidean_metric(query, proto) / self.args.temperature
# transform for all instances in the task
if mode == 'train':
aux_task = torch.cat([
support.reshape(self.args.shot, -1, support.shape[-1]),
query.reshape(self.args.query, -1, support.shape[-1])
], 0) # (K+Kq) x N x d
aux_task = aux_task.permute([1, 0, 2])
aux_emb = self.slf_attn(aux_task, aux_task,
aux_task) # N x (K+Kq) x d
# compute class mean
aux_center = torch.mean(aux_emb, 1) # N x d
logitis2 = euclidean_metric(
aux_task.permute([1, 0, 2]).view(-1, self.z_dim),
aux_center) / self.args.temperature2
return logitis, logitis2
else:
return logitis
def main():
args = parser.parse_args()
if not args.out_name:
args.out_name = osp.basename(osp.dirname(args.data_path))
pprint(vars(args))
if args.load_checkpoint and os.path.isdir(osp.join(args.logs_dir, args.out_name)):
features_dict,labels_dict,dataset = load_checkpoint(args)
else:
print("start new process")
print(args.load_checkpoint,os.path.isdir(osp.join(args.logs_dir, args.out_name)))
print("model init...")
model_weight = torch.load('./pretrained/modelEncoder_Ness_MINI_ProtoNet_MINI_5shot_10way_max_acc.pth')
model = AmdimNet(ndf=args.ndf, n_rkhs=args.rkhs, n_depth=args.nd)
model_dict = model.state_dict()
pretrained_dict = model_weight['model']
pretrained_dict = {k.replace('module.', ''): v for k, v in pretrained_dict.items() if k.replace('module.', '') in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model.eval()
model = model.cuda()
print("dataset init...")
dataset = Dataset(args)
dataloader = DataLoader(dataset,batch_size=args.batch_size,pin_memory=True)
features_dict = {}; labels_dict = {}
print("extracting features")
with torch.no_grad():
for i,batch in tqdm(enumerate(dataloader,1)):
transformed_imgs, paths, labels = batch
transformed_imgs = transformed_imgs.cuda()
features_batch = model(transformed_imgs)
for path,feature,label in zip(paths,features_batch,labels):
features_dict[path] = feature.cpu()
labels_dict[path] = label if label else None # if from query: label == ''
save_checkpoint(args,features_dict,labels_dict,dataset)
print("saving checkpoints log complete")
support_features = torch.stack([features_dict[path] for path in dataset.support ])
query_features = torch.stack([features_dict[path] for path in dataset.query])
print(torch.tensor(dataset.class_lens))
indices = torch.cumsum(torch.tensor(dataset.class_lens),dim=0)
indices = torch.cat([torch.tensor([0]),indices])
mean_support_featers = torch.stack([torch.mean(support_features[ indices[i]:indices[i+1] ],dim=0) for i in range(len(indices)-1)])
logits = euclidean_metric(query_features, mean_support_featers) / args.temperature # distance here is negative --> similarity
print(logits.shape)
return logits,dataset
def forward(self, data_shot, data_query):
proto = self.encoder(data_shot)
proto = proto.reshape(self.args.shot, self.args.way, -1).mean(dim=0)
logits = euclidean_metric(self.encoder(data_query),
proto) / self.args.temperature
return logits