From aae65e6ed81f8ddba76b42dde05f7ccf6516cbc3 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 1 Jan 2021 11:59:54 -0700 Subject: [PATCH] Mods to byol_resnet_playground for large batches --- codes/scripts/byol_resnet_playground.py | 74 +++++++++---------------- codes/scripts/extract_square_images.py | 11 ++-- 2 files changed, 31 insertions(+), 54 deletions(-) diff --git a/codes/scripts/byol_resnet_playground.py b/codes/scripts/byol_resnet_playground.py index 77e82e8a..a5491fd5 100644 --- a/codes/scripts/byol_resnet_playground.py +++ b/codes/scripts/byol_resnet_playground.py @@ -26,7 +26,7 @@ from utils.options import dict_to_nonedict def structural_euc_dist(x, y): diff = torch.square(x - y) sum = torch.sum(diff, dim=-1) - return torch.mean(torch.sqrt(sum)) + return torch.sqrt(sum) def cosine_similarity(x, y): @@ -87,43 +87,15 @@ def register_hook(net, layer_name): layer.register_forward_hook(_hook) -def create_latent_database(model, model_index=0): - batch_size = 32 - num_workers = 1 - output_path = '../../results/byol_resnet_latents/' - - os.makedirs(output_path, exist_ok=True) - dataloader = get_image_folder_dataloader(batch_size, num_workers) - id = 0 - dict_count = 1 - latent_dict = {} - all_paths = [] - for batch in tqdm(dataloader): - hq = batch['hq'].to('cuda') - latent = model(hq)[model_index] # BYOL trainer only trains the '4' output, which is indexed at [1]. Confusing. - for b in range(latent.shape[0]): - im_path = batch['HQ_path'][b] - all_paths.append(im_path) - latent_dict[id] = latent[b].detach().cpu() - if (id+1) % 1000 == 0: - print("Saving checkpoint..") - torch.save(latent_dict, os.path.join(output_path, "latent_dict_%i.pth" % (dict_count,))) - latent_dict = {} - torch.save(all_paths, os.path.join(output_path, "all_paths.pth")) - dict_count += 1 - id += 1 - - - def get_latent_for_img(model, img): img_t = ToTensor()(Image.open(img)).to('cuda').unsqueeze(0) _, _, h, w = img_t.shape # Center crop img_t and resize to 224. d = min(h, w) dh, dw = (h-d)//2, (w-d)//2 - if dh == 0: + if dw != 0: img_t = img_t[:, :, :, dw:-dw] - else: + elif dh != 0: img_t = img_t[:, :, dh:-dh, :] img_t = torch.nn.functional.interpolate(img_t, size=(224, 224), mode="area") model(img_t) @@ -134,36 +106,42 @@ def get_latent_for_img(model, img): def find_similar_latents(model, compare_fn=structural_euc_dist): global layer_hooked_value - img = 'F:\\4k6k\\datasets\\ns_images\\adrianna\\analyze\\analyze_xx\\yui_xx.jpg' + img = 'F:\\4k6k\\datasets\\ns_images\\adrianna\\analyze\\analyze_xx\\poon.jpg' #img = 'F:\\4k6k\\datasets\\ns_images\\adrianna\\analyze\\analyze_xx\\nicky_xx.jpg' output_path = '../../results/byol_resnet_similars' os.makedirs(output_path, exist_ok=True) - imglatent = get_latent_for_img(model, img) - _, c, h, w = imglatent.shape + imglatent = get_latent_for_img(model, img).squeeze().unsqueeze(0) + _, c = imglatent.shape - batch_size = 32 - num_workers = 1 + batch_size = 128 + num_workers = 8 dataloader = get_image_folder_dataloader(batch_size, num_workers) id = 0 + output_batch = 1 results = [] + result_paths = [] for batch in tqdm(dataloader): hq = batch['hq'].to('cuda') model(hq) - latent = layer_hooked_value - for b in range(latent.shape[0]): - im_path = batch['HQ_path'][b] - results.append((im_path, compare_fn(imglatent, latent[b].unsqueeze(0)).item())) - id += 1 - if id > 2000: - break - results.sort(key=lambda x: x[1]) - for i in range(50): - mag = results[i][1] - shutil.copy(results[i][0], os.path.join(output_path, f'{i}_{mag}.jpg')) + latent = layer_hooked_value.clone().squeeze() + compared = compare_fn(imglatent.repeat(latent.shape[0], 1), latent) + results.append(compared.cpu()) + result_paths.extend(batch['HQ_path']) + id += batch_size + if id > 10000: + k = 500 + results = torch.cat(results, dim=0) + vals, inds = torch.topk(results, k, largest=False) + for i in inds: + mag = int(results[i].item() * 1000) + shutil.copy(result_paths[i], os.path.join(output_path, f'{mag:05}_{output_batch}_{i}.jpg')) + results = [] + result_paths = [] + id = 0 if __name__ == '__main__': - pretrained_path = '../../experiments/resnet_byol_diffframe_69k.pth' + pretrained_path = '../../experiments/resnet_byol_diffframe_85k.pth' model = resnet50(pretrained=False).to('cuda') sd = torch.load(pretrained_path) resnet_sd = {} diff --git a/codes/scripts/extract_square_images.py b/codes/scripts/extract_square_images.py index 848da7aa..bf0777bf 100644 --- a/codes/scripts/extract_square_images.py +++ b/codes/scripts/extract_square_images.py @@ -19,9 +19,9 @@ def main(): # compression time. If read raw images during training, use 0 for faster IO speed. opt['dest'] = 'file' - opt['input_folder'] = ['F:\\4k6k\\datasets\\ns_images\\imagesets\\imgset4'] - opt['save_folder'] = 'F:\\4k6k\\datasets\\ns_images\\256_unsupervised_new' - opt['imgsize'] = 256 + opt['input_folder'] = ['F:\\4k6k\\datasets\\ns_images\\vixen\\vix_cropped'] + opt['save_folder'] = 'F:\\4k6k\\datasets\\ns_images\\video_512_cropped' + opt['imgsize'] = 512 #opt['bottom_crop'] = 120 save_folder = opt['save_folder'] @@ -45,7 +45,7 @@ class TiledDataset(data.Dataset): def get(self, index): path = self.images[index] basename = osp.basename(path) - img = data_util.read_img(None, path) + img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # Greyscale not supported. if img is None: @@ -62,7 +62,7 @@ class TiledDataset(data.Dataset): h, w, c = img.shape # Uncomment to filter any image that doesnt meet a threshold size. - if min(h,w) < 256: + if min(h,w) < 512: print("Skipping due to threshold") return None @@ -71,7 +71,6 @@ class TiledDataset(data.Dataset): # Crop the image so that only the center is left, since this is often the most salient part of the image. img = img[(h - dim) // 2:dim + (h - dim) // 2, (w - dim) // 2:dim + (w - dim) // 2, :] img = cv2.resize(img, (self.opt['imgsize'], self.opt['imgsize']), interpolation=cv2.INTER_AREA) - cv2.imwrite(osp.join(self.opt['save_folder'], basename + ".jpg"), img, [cv2.IMWRITE_JPEG_QUALITY, self.opt['compression_level']]) return None