<jupyter_start><jupyter_text>PKU-sketech dataset<jupyter_code>from config import Config_pku
from engine.pku_market_trainer import *
pku_args = Config_pku()
train_set, test_color_set, test_sketch_set = make_data_loader_pku(pku_args)

n_pku_clc = len(np.unique(train_set.train_data['sketch_labels'])) # number of person in training set
n_test_color = len(test_color_set.color_labels) #number of instance in test color set
n_test_sketch = len(test_sketch_set.sketch_labels) #number of instance in test sketch set

pku_train_loader = data.DataLoader(train_set, batch_size=pku_args.batch_size, shuffle=True, num_workers=pku_args.num_workers, drop_last=True)
pku_test_color_loader = data.DataLoader(test_color_set, batch_size=pku_args.test_batch, shuffle=False, num_workers=pku_args.num_workers)
pku_test_sketch_loader = data.DataLoader(test_sketch_set, batch_size=pku_args.test_batch, shuffle=False, num_workers=pku_args.num_workers)

test_color_labels = test_color_set.color_labels
test_sketch_labels = test_sketch_set.sketch_labels

nquery = n_test_sketch
ngall = n_test_color

query_loader = pku_test_sketch_loader
gall_loader = pku_test_color_loader 
query_label = test_sketch_labels
gall_label = test_color_labels
Embed_net_1 = Baseline(pretrain_choice=None)
Embed_net_2 = Baseline(pretrain_choice=None)
Classify_net = C_net(2048,n_pku_clc,0.7)
A_net = Attribute_net(dim=2048, n_att=8)

Embed_net_1.to(device)
Embed_net_2.to(device)
Classify_net.to(device)
A_net.to(device)

checkpoint = torch.load(pku_args.model_path+'pku_best_96.t')
Embed_net_1.load_state_dict(checkpoint['Embed_net_1'])
Embed_net_2.load_state_dict(checkpoint['Embed_net_2']) 
Classify_net.load_state_dict(checkpoint['Classify_net']) 

Embed_net_1.eval()
Embed_net_2.eval()
Classify_net.eval()
print ('Extracting Query Feature...')
ptr = 0
query_feat = np.zeros((nquery, 2048))
with torch.no_grad():
    for batch_idx, (img, label ) in enumerate(query_loader):
        batch_num = img.size(0)
        img = img.cuda()
        _,feat = Embed_net_2(img)
        feat = Classify_net(feat)
        feat = F.normalize(feat, p=2, dim=1) # normorlize features
        query_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()
        ptr = ptr + batch_num

        
print ('Extracting Gallery Feature...')
ptr = 0
gall_feat = np.zeros((ngall, 2048))
with torch.no_grad():
    for batch_idx, (img, label ) in enumerate(gall_loader):
        batch_num = img.size(0)
        img = img.cuda()
        _,feat = Embed_net_1(img)
        feat = Classify_net(feat)
        feat = F.normalize(feat, p=2, dim=1) # normorlize features
        gall_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()
        ptr = ptr + batch_num 
# compute the similarity
distmat  = np.matmul(query_feat, np.transpose(gall_feat))
cmc, mAP = eval_pku(-distmat, query_label, gall_label)
print(f'CMC from 1 to 20:{cmc}')
print(f'mAP:{mAP}')<jupyter_output>CMC from 1 to 20:[0.96 0.98 0.98 0.98 1.   1.   1.   1.   1.   1.   1.   1.   1.   1.
 1.   1.   1.   1.   1.   1.  ]
mAP:0.9430336134453782<jupyter_text>Shoe<jupyter_code>from config import Config_sbir
from engine.zap_sbir_trainer import *
sbir_args = Config_sbir()
train_set, test_color_set, test_sketch_set = make_data_loader_sbir(sbir_args)

n_sbir_clc = len(np.unique(train_set.train_data['labels'])) 
n_test_color = len(test_color_set.color_labels) 
n_test_sketch = len(test_sketch_set.sketch_labels) 


sbir_train_loader = data.DataLoader(train_set, batch_size=sbir_args.batch_size, shuffle=True, num_workers=sbir_args.num_workers, drop_last=False)
sbir_test_color_loader = data.DataLoader(test_color_set, batch_size=sbir_args.test_batch, shuffle=False, num_workers=sbir_args.num_workers)
sbir_test_sketch_loader = data.DataLoader(test_sketch_set, batch_size=sbir_args.test_batch, shuffle=False, num_workers=sbir_args.num_workers)

test_color_labels = np.array(test_color_set.color_labels)
test_sketch_labels = np.array(test_sketch_set.sketch_labels)

nquery = n_test_sketch
ngall = n_test_color
query_loader = sbir_test_sketch_loader
gall_loader = sbir_test_color_loader 
query_label = test_sketch_labels
gall_label = test_color_labels
Embed_net_1 = Baseline(pretrain_choice=None)
Embed_net_2 = Baseline(pretrain_choice=None)
Classify_net = C_net(sbir_args.low_dim,n_sbir_clc, 0.7)
A_net = Attribute_net(dim=sbir_args.low_dim, n_att=sbir_args.num_att)
D_net = Domain_net(dim=sbir_args.low_dim)

Embed_net_1.to(device)
Embed_net_2.to(device)
Classify_net.to(device)
A_net.to(device)

Embed_net_1 = nn.DataParallel(Embed_net_1)
Embed_net_2 = nn.DataParallel(Embed_net_2)
Classify_net = nn.DataParallel(Classify_net)
A_net = nn.DataParallel(A_net)

checkpoint = torch.load(sbir_args.model_path+'sbir_best_69.t')
Embed_net_1.load_state_dict(checkpoint['Embed_net_1'])
Embed_net_2.load_state_dict(checkpoint['Embed_net_2'])   

Embed_net_1.eval()
Embed_net_2.eval()
print ('Extracting Query Feature...')
ptr = 0
query_feat = np.zeros((nquery, 2048))
with torch.no_grad():
    for batch_idx, (img, label ) in enumerate(query_loader):
        batch_num = img.size(0)
        img = img.cuda()
        _,feat = Embed_net_2(img)
        feat = F.normalize(feat, p=2, dim=1) # normorlize features
        query_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()
        attributes = A_net(feat)
        ptr = ptr + batch_num         
print ('Extracting Gallery Feature...')
ptr = 0
gall_feat = np.zeros((ngall, 2048))

with torch.no_grad():
    for batch_idx, (img, label ) in enumerate(gall_loader):
        batch_num = img.size(0)
        img = img.cuda()
        _,feat = Embed_net_1(img)
        feat = F.normalize(feat, p=2, dim=1) # normorlize features
        gall_feat[ptr:ptr+batch_num,: ] = feat.detach().cpu().numpy()
        ptr = ptr + batch_num 

# compute the similarity
distmat  = np.matmul(query_feat, np.transpose(gall_feat))

from eval_metrics import eval_sbir
cmc, mAP = eval_sbir(-distmat, query_label, gall_label)
print(f'CMC from 1 to 20:{cmc}')
print(f'mAP:{mAP}')<jupyter_output>CMC from 1 to 20:[0.6869565  0.79130435 0.8521739  0.87826085 0.90434784 0.93043476
 0.9478261  0.9478261  0.9478261  0.9478261  0.95652175 0.95652175
 0.95652175 0.9652174  0.9652174  0.9652174  0.9652174  0.9652174
 0.9652174  0.9652174 ]
mAP:0.7802725383916311<jupyter_text>transfer weigh from parrall to single gpu<jupyter_code>Embed_net_1 = Baseline(pretrain_choice=None)
Embed_net_2 = Baseline(pretrain_choice=None)
Classify_net = C_net(sbir_args.low_dim,n_sbir_clc, 0.7)
A_net = Attribute_net(dim=sbir_args.low_dim, n_att=sbir_args.num_att)
D_net = Domain_net(dim=sbir_args.low_dim)

Embed_net_1 = nn.DataParallel(Embed_net_1)
Embed_net_2 = nn.DataParallel(Embed_net_2)
Classify_net = nn.DataParallel(Classify_net)
A_net = nn.DataParallel(A_net)
D_net = nn.DataParallel(D_net) 

checkpoint = torch.load(sbir_args.model_path+'zap_p.t')
Embed_net_1.load_state_dict(checkpoint['Embed_net_1'])
Embed_net_2.load_state_dict(checkpoint['Embed_net_2'])   
A_net.load_state_dict(checkpoint['A_net'])
      
state = {
    'Embed_net_1': Embed_net_1.module.state_dict(),
    'Embed_net_2': Embed_net_2.module.state_dict(),
    'Classify_net': Classify_net.module.state_dict(),
    'A_net': A_net.module.state_dict(),
}
torch.save(state, 'zap.t')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11556960285832074, 'fraction_numerical': 0.04589803490449361, 'mean_word_length': 5.544964028776978, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '221872', 'license': 'MIT', 'n_tokens_mistral': 3131, 'n_tokens_neox': 2909, 'n_words': 444, 'path': 'Testing.ipynb', 'repo_name': 'fandulu/IHDA', 'size': '7277'}