Spaces:
Sleeping
Sleeping
| import os | |
| import os.path as osp | |
| from copy import deepcopy | |
| from collections import OrderedDict | |
| import glob | |
| from datetime import datetime | |
| import random | |
| import copy | |
| import imageio | |
| import torch | |
| # import clip | |
| import torchvision.transforms.functional as tvf | |
| import video3d.utils.meters as meters | |
| import video3d.utils.misc as misc | |
| # from video3d.dataloaders import get_image_loader | |
| from video3d.dataloaders_ddp import get_sequence_loader_ddp, get_sequence_loader_quadrupeds, get_test_loader_quadrupeds | |
| from . import discriminator_architecture | |
| def sample_frames(batch, num_sample_frames, iteration, stride=1): | |
| ## window slicing sampling | |
| images, masks, flows, bboxs, bg_image, seq_idx, frame_idx = batch | |
| num_seqs, total_num_frames = images.shape[:2] | |
| # start_frame_idx = iteration % (total_num_frames - num_sample_frames +1) | |
| ## forward and backward | |
| num_windows = total_num_frames - num_sample_frames +1 | |
| start_frame_idx = (iteration * stride) % (2*num_windows) | |
| ## x' = (2n-1)/2 - |(2n-1)/2 - x| : 0,1,2,3,4,5 -> 0,1,2,2,1,0 | |
| mid_val = (2*num_windows -1) /2 | |
| start_frame_idx = int(mid_val - abs(mid_val -start_frame_idx)) | |
| new_batch = images[:, start_frame_idx:start_frame_idx+num_sample_frames], \ | |
| masks[:, start_frame_idx:start_frame_idx+num_sample_frames], \ | |
| flows[:, start_frame_idx:start_frame_idx+num_sample_frames-1], \ | |
| bboxs[:, start_frame_idx:start_frame_idx+num_sample_frames], \ | |
| bg_image, \ | |
| seq_idx, \ | |
| frame_idx[:, start_frame_idx:start_frame_idx+num_sample_frames] | |
| return new_batch | |
| def indefinite_generator(loader): | |
| while True: | |
| for x in loader: | |
| yield x | |
| def indefinite_generator_from_list(loaders): | |
| while True: | |
| random_idx = random.randint(0, len(loaders)-1) | |
| for x in loaders[random_idx]: | |
| yield x | |
| break | |
| def get_optimizer(model, lr=0.0001, betas=(0.9, 0.999), weight_decay=0): | |
| return torch.optim.Adam( | |
| filter(lambda p: p.requires_grad, model.parameters()), | |
| lr=lr, betas=betas, weight_decay=weight_decay) | |
| class Fewshot_Trainer: | |
| def __init__(self, cfgs, model): | |
| # only now supports one gpu | |
| self.cfgs = cfgs | |
| # here should be the one gpu ddp setting | |
| self.rank = cfgs.get('rank', 0) | |
| self.world_size = cfgs.get('world_size', 1) | |
| self.use_ddp = cfgs.get('use_ddp', True) | |
| self.device = cfgs.get('device', 'cpu') | |
| self.num_epochs = cfgs.get('num_epochs', 1) | |
| self.lr = cfgs.get('few_shot_lr', 1e-4) | |
| self.dataset = 'image' | |
| self.metrics_trace = meters.MetricsTrace() | |
| self.make_metrics = lambda m=None: meters.StandardMetrics(m) | |
| self.archive_code = cfgs.get('archive_code', True) | |
| self.batch_size = cfgs.get('batch_size', 64) | |
| self.in_image_size = cfgs.get('in_image_size', 256) | |
| self.out_image_size = cfgs.get('out_image_size', 256) | |
| self.num_workers = cfgs.get('num_workers', 4) | |
| self.checkpoint_dir = cfgs.get('checkpoint_dir', 'results') | |
| misc.xmkdir(self.checkpoint_dir) | |
| self.few_shot_resume = cfgs.get('few_shot_resume', False) | |
| self.save_checkpoint_freq = cfgs.get('save_checkpoint_freq', 1) | |
| self.keep_num_checkpoint = cfgs.get('keep_num_checkpoint', 2) # -1 for keeping all checkpoints | |
| self.few_shot_data_dir = cfgs.get('few_shot_data_dir', None) | |
| assert self.few_shot_data_dir is not None | |
| # in case we add more data source | |
| if isinstance(self.few_shot_data_dir, list): | |
| self.few_shot_data_dir_more = self.few_shot_data_dir[1:] | |
| self.few_shot_data_dir = self.few_shot_data_dir[0] | |
| else: | |
| self.few_shot_data_dir_more = None | |
| assert "data_resize_update" in self.few_shot_data_dir # TODO: a hack way to make sure not using wrong data, needs to remove | |
| self.few_shot_categories, self.few_shot_categories_paths = self.parse_few_shot_categories(self.few_shot_data_dir, self.few_shot_data_dir_more) | |
| # if we need test categories, we pop it from self.few_shot_categories and self.few_shot_categories_path | |
| # the test category needs to be category from few-shot, and we're using bs=1 on them, no need for back views enhancement (for now, use back view images, but don't duplicate them) | |
| self.test_category_num = cfgs.get('few_shot_test_category_num', 0) | |
| self.test_category_names = cfgs.get('few_shot_test_category_names', None) | |
| if self.test_category_num > 0: | |
| # if we have valid test_category names, then use them, the number doesn't need to be equal | |
| if self.test_category_names is not None: | |
| test_cats = self.test_category_names | |
| else: | |
| test_cats = list(self.few_shot_categories_paths.keys())[-(self.test_category_num):] | |
| test_categories_paths = {} | |
| for test_cat in test_cats: | |
| test_categories_paths.update({test_cat: self.few_shot_categories_paths[test_cat]}) | |
| assert test_cat in self.few_shot_categories | |
| self.few_shot_categories.remove(test_cat) | |
| self.few_shot_categories_paths.pop(test_cat) | |
| self.test_categories_paths = test_categories_paths | |
| else: | |
| self.test_categories_paths = None | |
| # also load the original 7 categories | |
| self.original_train_data_path = cfgs.get('train_data_dir', None) | |
| self.original_val_data_path = cfgs.get('val_data_dir', None) | |
| self.original_categories = [] | |
| self.original_categories_paths = self.original_train_data_path | |
| for k, v in self.original_train_data_path.items(): | |
| self.original_categories.append(k) | |
| self.categories = self.original_categories + self.few_shot_categories | |
| self.categories_paths = self.original_train_data_path.copy() | |
| self.categories_paths.update(self.few_shot_categories_paths) | |
| print(f'Using {len(self.categories)} cateogires: ', self.categories) | |
| # initialize new things | |
| # self.original_classes_num = cfgs.get('few_shot_original_classes_num', 7) | |
| self.original_classes_num = len(self.original_categories) | |
| self.new_classes_num = len(self.categories) - self.original_classes_num | |
| self.combine_dataset = cfgs.get('combine_dataset', False) | |
| assert self.combine_dataset, "we should use combine dataset, it's up to date" | |
| if self.combine_dataset: | |
| self.train_loader, self.val_loader, self.test_loader = self.get_data_loaders_quadrupeds(self.cfgs, self.batch_size, self.num_workers, self.in_image_size, self.out_image_size) | |
| else: | |
| self.train_loader_few_shot, self.val_loader_few_shot = self.get_data_loaders_few_shot(self.cfgs, self.batch_size, self.num_workers, self.in_image_size, self.out_image_size) | |
| self.train_loader_original, self.val_loader_original = self.get_data_loaders_original(self.cfgs, self.batch_size, self.num_workers, self.in_image_size, self.out_image_size) | |
| self.train_loader = self.train_loader_original + self.train_loader_few_shot | |
| if self.val_loader_few_shot is not None and self.val_loader_original is not None: | |
| self.val_loader = self.val_loader_original + self.val_loader_few_shot | |
| self.num_iterations = cfgs.get('num_iterations', 0) | |
| if self.num_iterations != 0: | |
| self.use_total_iterations = True | |
| else: | |
| self.use_total_iterations = False | |
| if self.use_total_iterations: | |
| # reset the epoch related cfgs | |
| dataloader_length = max([len(loader) for loader in self.train_loader]) * len(self.train_loader) | |
| print("Total length of data loader is: ", dataloader_length) | |
| total_epoch = int(self.num_iterations / dataloader_length) + 1 | |
| print(f'run for {total_epoch} epochs') | |
| print('is_main_process()?', misc.is_main_process()) | |
| for k, v in cfgs.items(): | |
| if 'epoch' in k: | |
| # if isinstance(v, list): | |
| # new_v = [int(total_epoch * x / 120) + 1 for x in v] | |
| # cfgs[k] = new_v | |
| # elif isinstance(v, int): | |
| # new_v = int(total_epoch * v / 120) + 1 | |
| # cfgs[k] = new_v | |
| # a better transformation | |
| if isinstance(v, int): | |
| # use the floor int | |
| new_v = int(total_epoch * v / 120) | |
| cfgs[k] = new_v | |
| elif isinstance(v, list): | |
| if v[0] == v[1]: | |
| # if the values in v are the same, then we use both the floor value | |
| new_v = [int(total_epoch * x / 120) for x in v] | |
| else: | |
| # if the values are not the same, make the first using floor value and others using ceil value | |
| new_v = [int(total_epoch * x / 120) + 1 for x in v] | |
| new_v[0] = new_v[0] - 1 | |
| cfgs[k] = new_v | |
| else: | |
| continue | |
| self.num_epochs = total_epoch | |
| self.cub_start_epoch = cfgs.get('cub_start_epoch', 0) | |
| self.cfgs = cfgs | |
| # the model is with nothing now | |
| self.model = model(cfgs) | |
| self.metrics_trace = meters.MetricsTrace() | |
| self.make_metrics = lambda m=None: meters.StandardMetrics(m) | |
| self.use_logger = True | |
| self.log_freq_images = cfgs.get('log_freq_images', 1000) | |
| self.log_train_images = cfgs.get('log_train_images', False) | |
| self.log_freq_losses = cfgs.get('log_freq_losses', 100) | |
| self.save_result_freq = cfgs.get('save_result_freq', None) | |
| self.train_result_dir = osp.join(self.checkpoint_dir, 'results') | |
| self.fix_viz_batch = cfgs.get('fix_viz_batch', False) | |
| self.visualize_validation = cfgs.get('visualize_validation', False) | |
| # self.visualize_validation = False | |
| self.iteration_save = cfgs.get('few_shot_iteration_save', False) | |
| self.iteration_save_freq = cfgs.get('few_shot_iteration_save_freq', 2000) | |
| self.enable_memory_bank = cfgs.get('enable_memory_bank', False) | |
| if self.enable_memory_bank: | |
| self.memory_bank_dim = 128 | |
| self.memory_bank_size = cfgs.get('memory_bank_size', 60) | |
| self.memory_bank_topk = cfgs.get('memory_bank_topk', 10) | |
| # assert self.memory_bank_topk < self.memory_bank_size | |
| assert self.memory_bank_topk <= self.memory_bank_size | |
| self.memory_retrieve = cfgs.get('memory_retrieve', 'cos-linear') | |
| self.memory_bank_init = cfgs.get('memory_bank_init', 'random') | |
| if self.memory_bank_init == 'copy': | |
| # use trained 7 embeddings to initialize | |
| num_piece = self.memory_bank_size // self.original_classes_num | |
| num_left = self.memory_bank_size - num_piece * self.original_classes_num | |
| tmp_1 = torch.empty_like(self.model.netPrior.classes_vectors) | |
| tmp_1 = tmp_1.copy_(self.model.netPrior.classes_vectors) | |
| tmp_1 = tmp_1.unsqueeze(0).repeat(num_piece, 1, 1) | |
| tmp_1 = tmp_1.reshape(tmp_1.shape[0] * tmp_1.shape[1], tmp_1.shape[-1]) | |
| if num_left > 0: | |
| tmp_2 = torch.empty_like(self.model.netPrior.classes_vectors) | |
| tmp_2 = tmp_2.copy_(self.model.netPrior.classes_vectors) | |
| tmp_2 = tmp_2[:num_left] | |
| tmp = torch.cat([tmp_1, tmp_2], dim=0) | |
| else: | |
| tmp = tmp_1 | |
| self.memory_bank = torch.nn.Parameter(tmp, requires_grad=True) | |
| elif self.memory_bank_init == 'random': | |
| self.memory_bank = torch.nn.Parameter(torch.nn.init.uniform_(torch.empty(self.memory_bank_size, self.memory_bank_dim), a=-0.05, b=0.05), requires_grad=True) | |
| else: | |
| raise NotImplementedError | |
| self.memory_encoder = cfgs.get('memory_encoder', 'DINO') # if DINO then just use the network encoder | |
| if self.memory_encoder == 'CLIP': | |
| self.clip_model, _ = clip.load('ViT-B/32', self.device) | |
| self.clip_model = self.clip_model.eval().requires_grad_(False) | |
| self.clip_mean = [0.48145466, 0.4578275, 0.40821073] | |
| self.clip_std = [0.26862954, 0.26130258, 0.27577711] | |
| self.clip_reso = 224 | |
| self.memory_bank_keys_dim = 512 | |
| elif self.memory_encoder == 'DINO': | |
| self.memory_bank_keys_dim = 384 | |
| else: | |
| raise NotImplementedError | |
| memory_bank_keys = torch.nn.init.uniform_(torch.empty(self.memory_bank_size, self.memory_bank_keys_dim), a=-0.05, b=0.05) | |
| self.memory_bank_keys = torch.nn.Parameter(memory_bank_keys, requires_grad=True) | |
| else: | |
| print("no memory bank, just use image embedding, this is only for one experiment!") | |
| self.memory_encoder = cfgs.get('memory_encoder', 'DINO') # if DINO then just use the network encoder | |
| if self.memory_encoder == 'CLIP': | |
| self.clip_model, _ = clip.load('ViT-B/32', self.device) | |
| self.clip_model = self.clip_model.eval().requires_grad_(False) | |
| self.clip_mean = [0.48145466, 0.4578275, 0.40821073] | |
| self.clip_std = [0.26862954, 0.26130258, 0.27577711] | |
| self.clip_reso = 224 | |
| self.memory_bank_keys_dim = 512 | |
| elif self.memory_encoder == 'DINO': | |
| self.memory_bank_keys_dim = 384 | |
| else: | |
| raise NotImplementedError | |
| self.prepare_model() | |
| def parse_few_shot_categories(self, data_dir, data_dir_more=None): | |
| # parse the categories data_dir | |
| few_shot_category_num = self.cfgs.get('few_shot_category_num', -1) | |
| assert few_shot_category_num != 0 | |
| categories = sorted(os.listdir(data_dir)) | |
| cnt = 0 | |
| category_names = [] | |
| category_names_paths = {} | |
| for category in categories: | |
| if osp.isdir(osp.join(self.few_shot_data_dir, category, 'train')): | |
| category_path = osp.join(self.few_shot_data_dir, category, 'train') | |
| category_names.append(category) | |
| category_names_paths.update({category: category_path}) | |
| cnt += 1 | |
| if few_shot_category_num > 0 and cnt >= few_shot_category_num: | |
| break | |
| # more data | |
| if data_dir_more is not None: | |
| for data_dir_one in data_dir_more: | |
| new_categories = os.listdir(data_dir_one) | |
| for new_category in new_categories: | |
| ''' | |
| if this category is not used before, add a new item | |
| if there is this category before, add the paths to original paths, | |
| if its a str, make it a list | |
| if its already a list, append it | |
| ''' | |
| if new_category not in category_names: | |
| #TODO: a hacky way here, if in new data there is category used in 7-cat, we just make it a new one | |
| if new_category in list(self.cfgs.get('train_data_dir', None).keys()): | |
| new_category = '_' + new_category | |
| category_names.append(new_category) | |
| category_names_paths.update({ | |
| new_category: osp.join(data_dir_one, new_category, 'train') | |
| }) | |
| else: | |
| old_category_path = category_names_paths[new_category] | |
| if isinstance(old_category_path, str): | |
| category_names_paths[new_category] = [ | |
| old_category_path, | |
| osp.join(data_dir_one, new_category, 'train') | |
| ] | |
| elif isinstance(old_category_path, list): | |
| old_category_path = old_category_path + [osp.join(data_dir_one, new_category, 'train')] | |
| category_names_paths[new_category] = old_category_path | |
| else: | |
| raise NotImplementedError | |
| # category_names = sorted(category_names) | |
| return category_names, category_names_paths | |
| def prepare_model(self): | |
| # here we prepare the model weights at outside | |
| # 1. load the pretrain weight | |
| # 2. initialize anything new, like new class vectors | |
| # 3. initialize new optimizer for chosen parameters | |
| assert self.original_classes_num == len(self.model.netPrior.category_id_map) | |
| # load pretrain | |
| # if not assigned few_shot_checkpoint_name, then skip this part | |
| if self.cfgs.get('few_shot_checkpoint_name', None) is not None: | |
| original_checkpoint_path = osp.join(self.checkpoint_dir, self.cfgs.get('few_shot_checkpoint_name', 'checkpoint060.pth')) | |
| assert osp.exists(original_checkpoint_path) | |
| print(f"Loading pre-trained checkpoint from {original_checkpoint_path}") | |
| cp = torch.load(original_checkpoint_path, map_location=self.device) | |
| # if using local-texture network in fine-tuning, the texture in previous pre-train ckpt is global | |
| # here we use a hack way, we just get rid of original texture ckpt | |
| if (self.cfgs.get('texture_way', None) is not None) or (self.cfgs.get('texture_act', 'relu') != 'relu'): | |
| new_netInstance_weights = {k: v for k, v in cp['netInstance'].items() if 'netTexture' not in k} | |
| #find the new texture weights | |
| texture_weights = self.model.netInstance.netTexture.state_dict() | |
| #add the new weights to the new model weights | |
| for k, v in texture_weights.items(): | |
| # for the overlapping part in netTexture, we also use them | |
| # if ('netTexture.' + k) in cp['netInstance'].keys(): | |
| # new_netInstance_weights['netTexture.' + k] = cp['netInstance']['netTexture.' + k] | |
| # else: | |
| # new_netInstance_weights['netTexture.' + k] = v | |
| new_netInstance_weights['netTexture.' + k] = v | |
| _ = cp.pop("netInstance") | |
| cp.update({"netInstance": new_netInstance_weights}) | |
| self.model.netInstance.load_state_dict(cp["netInstance"], strict=False) # For Deform | |
| # self.model.netInstance.load_state_dict(cp["netInstance"]) | |
| self.model.netPrior.load_state_dict(cp["netPrior"]) | |
| self.original_total_iter = cp["total_iter"] | |
| else: | |
| print("not load any pre-train weight, the iter will start from 0, make sure you set all the needed parameters") | |
| self.original_total_iter = 0 | |
| if not self.cfgs.get('disable_fewshot', False): | |
| for i, category in enumerate(self.few_shot_categories): | |
| category_id = self.original_classes_num + i | |
| self.model.netPrior.category_id_map.update({category: category_id}) | |
| few_shot_class_vector_init = self.cfgs.get('few_shot_class_vector_init', 'random') | |
| if few_shot_class_vector_init == 'random': | |
| tmp = torch.nn.init.uniform_(torch.empty(self.new_classes_num, self.model.netPrior.classes_vectors.shape[-1]), a=-0.05, b=0.05) | |
| tmp = tmp.to(self.model.netPrior.classes_vectors.device) | |
| self.model.netPrior.classes_vectors = torch.nn.Parameter(torch.cat([self.model.netPrior.classes_vectors, tmp], dim=0)) | |
| elif few_shot_class_vector_init == 'copy': | |
| num_7_cat_piece = self.new_classes_num // self.original_classes_num if self.new_classes_num > self.original_classes_num else 0 | |
| num_left = self.new_classes_num - num_7_cat_piece * self.original_classes_num | |
| if num_7_cat_piece > 0: | |
| tmp_1 = torch.empty_like(self.model.netPrior.classes_vectors) | |
| tmp_1 = tmp_1.copy_(self.model.netPrior.classes_vectors) | |
| tmp_1 = tmp_1.unsqueeze(0).repeat(num_7_cat_piece, 1, 1) | |
| tmp_1 = tmp_1.reshape(tmp_1.shape[0] * tmp_1.shape[1], tmp_1.shape[-1]) | |
| else: | |
| tmp_1 = None | |
| if num_left > 0: | |
| tmp_2 = torch.empty_like(self.model.netPrior.classes_vectors) | |
| tmp_2 = tmp_2.copy_(self.model.netPrior.classes_vectors) | |
| tmp_2 = tmp_2[:num_left] | |
| else: | |
| tmp_2 = None | |
| if tmp_1 != None and tmp_2 != None: | |
| tmp = torch.cat([tmp_1, tmp_2], dim=0) | |
| elif tmp_1 == None and tmp_2 != None: | |
| tmp = tmp_2 | |
| elif tmp_2 == None and tmp_1 != None: | |
| tmp = tmp_1 | |
| else: | |
| raise NotImplementedError | |
| tmp = tmp.to(self.model.netPrior.classes_vectors.device) | |
| self.model.netPrior.classes_vectors = torch.nn.Parameter(torch.cat([self.model.netPrior.classes_vectors, tmp], dim=0)) | |
| else: | |
| raise NotImplementedError | |
| else: | |
| print("disable few shot, not increasing embedding vectors") | |
| # initialize new optimizer | |
| optimize_rule = self.cfgs.get('few_shot_optimize', 'all') | |
| if optimize_rule == 'all': | |
| optimize_list = [ | |
| {'name': 'net_Prior', 'params': list(self.model.netPrior.parameters()), 'lr': self.lr * 10.}, | |
| {'name': 'net_Instance', 'params': list(self.model.netInstance.parameters()), 'lr': self.lr * 1.}, | |
| ] | |
| elif optimize_rule == 'only-emb': | |
| optimize_list = [ | |
| {'name': 'class_embeddings', 'params': list([self.model.netPrior.classes_vectors]), 'lr': self.lr * 10.} | |
| ] | |
| elif optimize_rule == 'emb-instance': | |
| optimize_list = [ | |
| {'name': 'class_embeddings', 'params': list([self.model.netPrior.classes_vectors]), 'lr': self.lr * 10.}, | |
| {'name': 'net_Instance', 'params': list(self.model.netInstance.parameters()), 'lr': self.lr * 1.}, | |
| ] | |
| elif optimize_rule == 'custom': | |
| optimize_list = [ | |
| {'name': 'net_Prior', 'params': list(self.model.netPrior.parameters()), 'lr': self.lr * 10.}, | |
| {'name': 'netEncoder', 'params': list(self.model.netInstance.netEncoder.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netTexture', 'params': list(self.model.netInstance.netTexture.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netPose', 'params': list(self.model.netInstance.netPose.parameters()), 'lr': self.lr * 0.01}, | |
| {'name': 'netArticulation', 'params': list(self.model.netInstance.netArticulation.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netLight', 'params': list(self.model.netInstance.netLight.parameters()), 'lr': self.lr * 1.} | |
| ] | |
| elif optimize_rule == 'custom-deform': | |
| optimize_list = [ | |
| {'name': 'net_Prior', 'params': list(self.model.netPrior.parameters()), 'lr': self.lr * 10.}, | |
| {'name': 'netEncoder', 'params': list(self.model.netInstance.netEncoder.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netTexture', 'params': list(self.model.netInstance.netTexture.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netPose', 'params': list(self.model.netInstance.netPose.parameters()), 'lr': self.lr * 0.01}, | |
| {'name': 'netArticulation', 'params': list(self.model.netInstance.netArticulation.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netLight', 'params': list(self.model.netInstance.netLight.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netDeform', 'params': list(self.model.netInstance.netDeform.parameters()), 'lr': self.lr * 1.} | |
| ] | |
| elif optimize_rule == 'texture': | |
| optimize_list = [ | |
| {'name': 'netTexture', 'params': list(self.model.netInstance.netTexture.parameters()), 'lr': self.lr * 1.} | |
| ] | |
| elif optimize_rule == 'texture-light': | |
| optimize_list = [ | |
| {'name': 'netTexture', 'params': list(self.model.netInstance.netTexture.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netLight', 'params': list(self.model.netInstance.netLight.parameters()), 'lr': self.lr * 1.} | |
| ] | |
| elif optimize_rule == 'exp': | |
| optimize_list = [ | |
| {'name': 'net_Prior', 'params': list(self.model.netPrior.parameters()), 'lr': self.lr * 10.}, | |
| {'name': 'netEncoder', 'params': list(self.model.netInstance.netEncoder.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netTexture', 'params': list(self.model.netInstance.netTexture.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netPose', 'params': list(self.model.netInstance.netPose.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netArticulation', 'params': list(self.model.netInstance.netArticulation.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netLight', 'params': list(self.model.netInstance.netLight.parameters()), 'lr': self.lr * 1.}, | |
| {'name': 'netDeform', 'params': list(self.model.netInstance.netDeform.parameters()), 'lr': self.lr * 1.} | |
| ] | |
| else: | |
| raise NotImplementedError | |
| if self.enable_memory_bank and optimize_rule != 'texture': | |
| optimize_bank_components = self.cfgs.get('few_shot_optimize_bank', 'all') | |
| if optimize_bank_components == 'value': | |
| optimize_list += [ | |
| {'name': 'memory_bank', 'params': list([self.memory_bank]), 'lr': self.lr * 10.} | |
| ] | |
| elif optimize_bank_components == 'key': | |
| optimize_list += [ | |
| {'name': 'memory_bank_keys', 'params': list([self.memory_bank_keys]), 'lr': self.lr * 10.} | |
| ] | |
| else: | |
| optimize_list += [ | |
| {'name': 'memory_bank', 'params': list([self.memory_bank]), 'lr': self.lr * 10.}, | |
| {'name': 'memory_bank_keys', 'params': list([self.memory_bank_keys]), 'lr': self.lr * 10.} | |
| ] | |
| if self.model.enable_vsd: | |
| optimize_list += [ | |
| {'name': 'lora', 'params': list(self.model.stable_diffusion.parameters()), 'lr': self.lr} | |
| ] | |
| # self.optimizerFewShot = torch.optim.Adam( | |
| # [ | |
| # # {'name': 'class_embeddings', 'params': list([self.model.netPrior.classes_vectors]), 'lr': self.lr * 1.}, | |
| # {'name': 'net_Prior', 'params': list(self.model.netPrior.parameters()), 'lr': self.lr * 10.}, | |
| # {'name': 'net_Instance', 'params': list(self.model.netInstance.parameters()), 'lr': self.lr * 1.}, | |
| # # {'name': 'net_articulation', 'params': list(self.model.netInstance.netArticulation.parameters()), 'lr': self.lr * 10.} | |
| # ], betas=(0.9, 0.99), eps=1e-15 | |
| # ) | |
| self.optimizerFewShot = torch.optim.Adam(optimize_list, betas=(0.9, 0.99), eps=1e-15) | |
| # if self.cfgs.get('texture_way', None) is not None and self.cfgs.get('gan_tex', False): | |
| if self.cfgs.get('gan_tex', False): | |
| self.optimizerDiscTex = torch.optim.Adam(filter(lambda p: p.requires_grad, self.model.discriminator_texture.parameters()), lr=self.lr, betas=(0.9, 0.99), eps=1e-15) | |
| def load_checkpoint(self, optim=True, checkpoint_name=None): | |
| # use to load the checkpoint of model and optimizer in the finetuning | |
| """Search the specified/latest checkpoint in checkpoint_dir and load the model and optimizer.""" | |
| if checkpoint_name is not None: | |
| checkpoint_path = osp.join(self.checkpoint_dir, checkpoint_name) | |
| else: | |
| checkpoints = sorted(glob.glob(osp.join(self.checkpoint_dir, '*.pth'))) | |
| if len(checkpoints) == 0: | |
| return 0, 0 | |
| checkpoint_path = checkpoints[-1] | |
| self.checkpoint_name = osp.basename(checkpoint_path) | |
| print(f"Loading checkpoint from {checkpoint_path}") | |
| cp = torch.load(checkpoint_path, map_location=self.device) | |
| self.model.load_model_state(cp) # the cp has netPrior and netInstance as keys | |
| if optim: | |
| try: | |
| self.optimizerFewShot.load_state_dict(cp['optimizerFewShot']) | |
| except: | |
| print('you should be using the local texture so dont need to load the previous optimizer') | |
| if self.enable_memory_bank: | |
| self.memory_bank_keys = cp['memory_bank_keys'] | |
| self.memory_bank = cp['memory_bank'] | |
| self.metrics_trace = cp['metrics_trace'] | |
| epoch = cp['epoch'] | |
| total_iter = cp['total_iter'] | |
| return epoch, total_iter | |
| def save_checkpoint(self, epoch, total_iter=0, optim=True, use_iter=False): | |
| """Save model, optimizer, and metrics state to a checkpoint in checkpoint_dir for the specified epoch.""" | |
| misc.xmkdir(self.checkpoint_dir) | |
| if use_iter: | |
| checkpoint_path = osp.join(self.checkpoint_dir, f'iter{total_iter:07}.pth') | |
| prefix = 'iter*.pth' | |
| else: | |
| checkpoint_path = osp.join(self.checkpoint_dir, f'checkpoint{epoch:03}.pth') | |
| prefix = 'checkpoint*.pth' | |
| state_dict = self.model.get_model_state() | |
| if optim: | |
| optimizer_state = {'optimizerFewShot': self.optimizerFewShot.state_dict()} | |
| state_dict = {**state_dict, **optimizer_state} | |
| state_dict['metrics_trace'] = self.metrics_trace | |
| state_dict['epoch'] = epoch | |
| state_dict['total_iter'] = total_iter | |
| if self.enable_memory_bank: | |
| state_dict['memory_bank_keys'] = self.memory_bank_keys | |
| state_dict['memory_bank'] = self.memory_bank | |
| print(f"Saving checkpoint to {checkpoint_path}") | |
| torch.save(state_dict, checkpoint_path) | |
| if self.keep_num_checkpoint > 0: | |
| self.clean_checkpoint(self.checkpoint_dir, keep_num=self.keep_num_checkpoint, prefix=prefix) | |
| def clean_checkpoint(self, checkpoint_dir, keep_num=2, prefix='checkpoint*.pth'): | |
| if keep_num > 0: | |
| names = list(sorted( | |
| glob.glob(os.path.join(checkpoint_dir, prefix)) | |
| )) | |
| if len(names) > keep_num: | |
| for name in names[:-keep_num]: | |
| print(f"Deleting obslete checkpoint file {name}") | |
| os.remove(name) | |
| def get_data_loaders_few_shot(self, cfgs, batch_size, num_workers, in_image_size, out_image_size): | |
| # support the train_data_loaders, and also an identical val_data_loader? | |
| train_loader = val_loader = None | |
| color_jitter_train = cfgs.get('color_jitter_train', None) | |
| color_jitter_val = cfgs.get('color_jitter_val', None) | |
| random_flip_train = cfgs.get('random_flip_train', False) | |
| data_loader_mode = cfgs.get('data_loader_mode', 'n_frame') | |
| num_sample_frames = cfgs.get('num_sample_frames', 2) | |
| shuffle_train_seqs = cfgs.get('shuffle_train_seqs', False) | |
| load_background = cfgs.get('background_mode', 'none') == 'background' | |
| rgb_suffix = cfgs.get('rgb_suffix', '.png') | |
| load_dino_feature = cfgs.get('load_dino_feature', False) | |
| dino_feature_dim = cfgs.get('dino_feature_dim', 64) | |
| get_loader_ddp = lambda **kwargs: get_sequence_loader_ddp( | |
| mode=data_loader_mode, | |
| batch_size=batch_size, | |
| num_workers=num_workers, | |
| in_image_size=in_image_size, | |
| out_image_size=out_image_size, | |
| num_sample_frames=num_sample_frames, | |
| load_background=load_background, | |
| rgb_suffix=rgb_suffix, | |
| load_dino_feature=load_dino_feature, | |
| dino_feature_dim=dino_feature_dim, | |
| flow_bool=0, | |
| **kwargs) | |
| print(f"Loading training data...") | |
| train_loader = get_loader_ddp(data_dir=[self.original_classes_num, self.few_shot_categories_paths], rank=self.rank, world_size=self.world_size, use_few_shot=True, shuffle=False, color_jitter=color_jitter_train, random_flip=random_flip_train) | |
| return train_loader, val_loader | |
| def get_data_loaders_original(self, cfgs, batch_size, num_workers, in_image_size, out_image_size): | |
| train_loader = val_loader = test_loader = None | |
| color_jitter_train = cfgs.get('color_jitter_train', None) | |
| color_jitter_val = cfgs.get('color_jitter_val', None) | |
| random_flip_train = cfgs.get('random_flip_train', False) | |
| data_loader_mode = cfgs.get('data_loader_mode', 'n_frame') | |
| skip_beginning = cfgs.get('skip_beginning', 4) | |
| skip_end = cfgs.get('skip_end', 4) | |
| num_sample_frames = cfgs.get('num_sample_frames', 2) | |
| min_seq_len = cfgs.get('min_seq_len', 10) | |
| max_seq_len = cfgs.get('max_seq_len', 10) | |
| debug_seq = cfgs.get('debug_seq', False) | |
| random_sample_train_frames = cfgs.get('random_sample_train_frames', False) | |
| shuffle_train_seqs = cfgs.get('shuffle_train_seqs', False) | |
| random_sample_val_frames = cfgs.get('random_sample_val_frames', False) | |
| load_background = cfgs.get('background_mode', 'none') == 'background' | |
| rgb_suffix = cfgs.get('rgb_suffix', '.png') | |
| load_dino_feature = cfgs.get('load_dino_feature', False) | |
| load_dino_cluster = cfgs.get('load_dino_cluster', False) | |
| dino_feature_dim = cfgs.get('dino_feature_dim', 64) | |
| get_loader_ddp = lambda **kwargs: get_sequence_loader_ddp( | |
| mode=data_loader_mode, | |
| batch_size=batch_size, | |
| num_workers=num_workers, | |
| in_image_size=in_image_size, | |
| out_image_size=out_image_size, | |
| debug_seq=debug_seq, | |
| skip_beginning=skip_beginning, | |
| skip_end=skip_end, | |
| num_sample_frames=num_sample_frames, | |
| min_seq_len=min_seq_len, | |
| max_seq_len=max_seq_len, | |
| load_background=load_background, | |
| rgb_suffix=rgb_suffix, | |
| load_dino_feature=load_dino_feature, | |
| load_dino_cluster=load_dino_cluster, | |
| dino_feature_dim=dino_feature_dim, | |
| flow_bool=0, | |
| **kwargs) | |
| # just the train now | |
| train_data_dir = self.original_categories_paths | |
| if isinstance(train_data_dir, dict): | |
| for data_path in train_data_dir.values(): | |
| assert osp.isdir(data_path), f"Training data directory does not exist: {data_path}" | |
| elif isinstance(train_data_dir, str): | |
| assert osp.isdir(train_data_dir), f"Training data directory does not exist: {train_data_dir}" | |
| else: | |
| raise ValueError("train_data_dir must be a string or a dict of strings") | |
| print(f"Loading training data...") | |
| # the train_data_dir is a dict and will go into the original dataset type | |
| train_loader = get_loader_ddp(data_dir=train_data_dir, rank=self.rank, world_size=self.world_size, is_validation=False, use_few_shot=False, random_sample=random_sample_train_frames, shuffle=shuffle_train_seqs, dense_sample=True, color_jitter=color_jitter_train, random_flip=random_flip_train) | |
| return train_loader, val_loader | |
| def get_data_loaders_quadrupeds(self, cfgs, batch_size, num_workers, in_image_size, out_image_size): | |
| train_loader = val_loader = test_loader = None | |
| color_jitter_train = cfgs.get('color_jitter_train', None) | |
| color_jitter_val = cfgs.get('color_jitter_val', None) | |
| random_flip_train = cfgs.get('random_flip_train', False) | |
| data_loader_mode = cfgs.get('data_loader_mode', 'n_frame') | |
| skip_beginning = cfgs.get('skip_beginning', 4) | |
| skip_end = cfgs.get('skip_end', 4) | |
| num_sample_frames = cfgs.get('num_sample_frames', 2) | |
| min_seq_len = cfgs.get('min_seq_len', 10) | |
| max_seq_len = cfgs.get('max_seq_len', 10) | |
| debug_seq = cfgs.get('debug_seq', False) | |
| random_sample_train_frames = cfgs.get('random_sample_train_frames', False) | |
| shuffle_train_seqs = cfgs.get('shuffle_train_seqs', False) | |
| random_sample_val_frames = cfgs.get('random_sample_val_frames', False) | |
| load_background = cfgs.get('background_mode', 'none') == 'background' | |
| rgb_suffix = cfgs.get('rgb_suffix', '.png') | |
| load_dino_feature = cfgs.get('load_dino_feature', False) | |
| load_dino_cluster = cfgs.get('load_dino_cluster', False) | |
| dino_feature_dim = cfgs.get('dino_feature_dim', 64) | |
| enhance_back_view = cfgs.get('enhance_back_view', False) | |
| enhance_back_view_path = cfgs.get('enhance_back_view_path', None) | |
| override_categories = cfgs.get('override_categories', None) | |
| disable_fewshot = cfgs.get('disable_fewshot', False) | |
| dataset_split_num = cfgs.get('dataset_split_num', -1) | |
| get_loader_ddp = lambda **kwargs: get_sequence_loader_quadrupeds( | |
| mode=data_loader_mode, | |
| num_workers=num_workers, | |
| in_image_size=in_image_size, | |
| out_image_size=out_image_size, | |
| debug_seq=debug_seq, | |
| skip_beginning=skip_beginning, | |
| skip_end=skip_end, | |
| num_sample_frames=num_sample_frames, | |
| min_seq_len=min_seq_len, | |
| max_seq_len=max_seq_len, | |
| load_background=load_background, | |
| rgb_suffix=rgb_suffix, | |
| load_dino_feature=load_dino_feature, | |
| load_dino_cluster=load_dino_cluster, | |
| dino_feature_dim=dino_feature_dim, | |
| flow_bool=0, | |
| enhance_back_view=enhance_back_view, | |
| enhance_back_view_path=enhance_back_view_path, | |
| override_categories=override_categories, | |
| disable_fewshot=disable_fewshot, | |
| dataset_split_num=dataset_split_num, | |
| **kwargs) | |
| # just the train now | |
| print(f"Loading training data...") | |
| val_image_num = cfgs.get('few_shot_val_image_num', 5) | |
| # the train_data_dir is a dict and will go into the original dataset type | |
| train_loader = get_loader_ddp(original_data_dirs=self.original_categories_paths, few_shot_data_dirs=self.few_shot_categories_paths, original_num=self.original_classes_num, few_shot_num=self.new_classes_num, rank=self.rank, world_size=self.world_size, batch_size=batch_size, is_validation=False, val_image_num=val_image_num, shuffle=shuffle_train_seqs, dense_sample=True, color_jitter=color_jitter_train, random_flip=random_flip_train) | |
| val_loader = get_loader_ddp(original_data_dirs=self.original_val_data_path, few_shot_data_dirs=self.few_shot_categories_paths, original_num=self.original_classes_num, few_shot_num=self.new_classes_num, rank=self.rank, world_size=self.world_size, batch_size=1, is_validation=True, val_image_num=val_image_num, shuffle=False, dense_sample=True, color_jitter=color_jitter_val, random_flip=False) | |
| if self.test_categories_paths is not None: | |
| get_test_loader_ddp = lambda **kwargs: get_test_loader_quadrupeds( | |
| mode=data_loader_mode, | |
| num_workers=num_workers, | |
| in_image_size=in_image_size, | |
| out_image_size=out_image_size, | |
| debug_seq=debug_seq, | |
| skip_beginning=skip_beginning, | |
| skip_end=skip_end, | |
| num_sample_frames=num_sample_frames, | |
| min_seq_len=min_seq_len, | |
| max_seq_len=max_seq_len, | |
| load_background=load_background, | |
| rgb_suffix=rgb_suffix, | |
| load_dino_feature=load_dino_feature, | |
| load_dino_cluster=load_dino_cluster, | |
| dino_feature_dim=dino_feature_dim, | |
| flow_bool=0, | |
| enhance_back_view=enhance_back_view, | |
| enhance_back_view_path=enhance_back_view_path, | |
| **kwargs) | |
| print(f"Loading testing data...") | |
| test_loader = get_test_loader_ddp(test_data_dirs=self.test_categories_paths, rank=self.rank, world_size=self.world_size, batch_size=1, is_validation=True, shuffle=False, dense_sample=True, color_jitter=color_jitter_val, random_flip=False) | |
| else: | |
| test_loader = None | |
| return train_loader, val_loader, test_loader | |
| def forward_frozen_ViT(self, images): | |
| # this part use the frozen pre-train ViT | |
| x = images | |
| with torch.no_grad(): | |
| b, c, h, w = x.shape | |
| self.model.netInstance.netEncoder._feats = [] | |
| self.model.netInstance.netEncoder._register_hooks([11], 'key') | |
| #self._register_hooks([11], 'token') | |
| x = self.model.netInstance.netEncoder.ViT.prepare_tokens(x) | |
| #x = self.ViT.prepare_tokens_with_masks(x) | |
| for blk in self.model.netInstance.netEncoder.ViT.blocks: | |
| x = blk(x) | |
| out = self.model.netInstance.netEncoder.ViT.norm(x) | |
| self.model.netInstance.netEncoder._unregister_hooks() | |
| ph, pw = h // self.model.netInstance.netEncoder.patch_size, w // self.model.netInstance.netEncoder.patch_size | |
| patch_out = out[:, 1:] # first is class token | |
| patch_out = patch_out.reshape(b, ph, pw, self.model.netInstance.netEncoder.vit_feat_dim).permute(0, 3, 1, 2) | |
| patch_key = self.model.netInstance.netEncoder._feats[0][:,:,1:] # B, num_heads, num_patches, dim | |
| patch_key = patch_key.permute(0, 1, 3, 2).reshape(b, self.model.netInstance.netEncoder.vit_feat_dim, ph, pw) | |
| global_feat = out[:, 0] | |
| return global_feat | |
| def forward_fix_embeddings(self, batch): | |
| images = batch[0] | |
| images = images.to(self.device) | |
| batch_size, num_frames, _, h0, w0 = images.shape | |
| images = images.reshape(batch_size*num_frames, *images.shape[2:]) # 0~1 | |
| if self.memory_encoder == 'DINO': | |
| images_in = images * 2 - 1 # rescale to (-1, 1) | |
| batch_features = self.forward_frozen_ViT(images_in) | |
| elif self.memory_encoder == 'CLIP': | |
| images_in = torch.nn.functional.interpolate(images, (self.clip_reso, self.clip_reso), mode='bilinear') | |
| images_in = tvf.normalize(images_in, self.clip_mean, self.clip_std) | |
| batch_features = self.clip_model.encode_image(images_in).float() | |
| else: | |
| raise NotImplementedError | |
| return batch_features | |
| def retrieve_memory_bank(self, batch_features, batch): | |
| batch_size = batch_features.shape[0] | |
| if self.memory_retrieve == 'cos-linear': | |
| query = torch.nn.functional.normalize(batch_features.unsqueeze(1), dim=-1) # [B, 1, d_k] | |
| key = torch.nn.functional.normalize(self.memory_bank_keys, dim=-1) # [size, d_k] | |
| key = key.transpose(1, 0).unsqueeze(0).repeat(batch_size, 1, 1).to(query.device) # [B, d_k, size] | |
| cos_dist = torch.bmm(query, key).squeeze(1) # [B, size], larger the more similar | |
| rank_idx = torch.sort(cos_dist, dim=-1, descending=True)[1][:, :self.memory_bank_topk] # [B, k] | |
| value = self.memory_bank.unsqueeze(0).repeat(batch_size, 1, 1).to(query.device) # [B, size, d_v] | |
| out = torch.gather(value, dim=1, index=rank_idx[..., None].repeat(1, 1, self.memory_bank_dim)) # [B, k, d_v] | |
| weights = torch.gather(cos_dist, dim=-1, index=rank_idx) # [B, k] | |
| weights = torch.nn.functional.normalize(weights, p=1.0, dim=-1).unsqueeze(-1).repeat(1, 1, self.memory_bank_dim) # [B, k, d_v] weights have been normalized | |
| out = weights * out | |
| out = torch.sum(out, dim=1) | |
| else: | |
| raise NotImplementedError | |
| batch_mean_out = torch.mean(out, dim=0) | |
| weight_aux = { | |
| 'weights': weights[:, :, 0], # [B, k], weights from large to small | |
| 'pick_idx': rank_idx, # [B, k] | |
| } | |
| return batch_mean_out, out, weight_aux | |
| def discriminator_texture_step(self): | |
| image_iv = self.model.record_image_iv | |
| image_rv = self.model.record_image_rv | |
| image_gt = self.model.record_image_gt | |
| self.model.record_image_iv = None | |
| self.model.record_image_rv = None | |
| self.model.record_image_gt = None | |
| image_iv = image_iv.requires_grad_(True) | |
| image_rv = image_rv.requires_grad_(True) | |
| image_gt = image_gt.requires_grad_(True) | |
| self.optimizerDiscTex.zero_grad() | |
| disc_loss_gt = 0.0 | |
| disc_loss_iv = 0.0 | |
| disc_loss_rv = 0.0 | |
| grad_penalty = 0.0 | |
| # for the gt image, it can only be in real or not | |
| if 'gt' in self.model.few_shot_gan_tex_real: | |
| d_gt = self.model.discriminator_texture(image_gt) | |
| disc_loss_gt += discriminator_architecture.bce_loss_target(d_gt, 1) | |
| if image_gt.requires_grad: | |
| grad_penalty_gt = 10. * discriminator_architecture.compute_grad2(d_gt, image_gt) | |
| disc_loss_gt += grad_penalty_gt | |
| grad_penalty += grad_penalty_gt | |
| # for the input view image, it can be in real or fake | |
| if 'iv' in self.model.few_shot_gan_tex_real: | |
| d_iv = self.model.discriminator_texture(image_iv) | |
| disc_loss_iv += discriminator_architecture.bce_loss_target(d_iv, 1) | |
| if image_iv.requires_grad: | |
| grad_penalty_iv = 10. * discriminator_architecture.compute_grad2(d_iv, image_iv) | |
| disc_loss_iv += grad_penalty_iv | |
| grad_penalty += grad_penalty_iv | |
| elif 'iv' in self.model.few_shot_gan_tex_fake: | |
| d_iv = self.model.discriminator_texture(image_iv) | |
| disc_loss_iv += discriminator_architecture.bce_loss_target(d_iv, 0) | |
| # for the random view image, it can only be in fake | |
| if 'rv' in self.model.few_shot_gan_tex_fake: | |
| d_rv = self.model.discriminator_texture(image_rv) | |
| disc_loss_rv += discriminator_architecture.bce_loss_target(d_rv, 0) | |
| all_loss = disc_loss_iv + disc_loss_rv + disc_loss_gt | |
| all_loss = all_loss * self.cfgs.get('gan_tex_loss_discriminator_weight', 0.1) | |
| self.discriminator_texture_loss = all_loss | |
| self.discriminator_texture_loss.backward() | |
| self.optimizerDiscTex.step() | |
| self.discriminator_texture_loss = 0. | |
| return { | |
| 'discriminator_loss': all_loss.detach(), | |
| 'discriminator_loss_iv': disc_loss_iv.detach(), | |
| 'discriminator_loss_rv': disc_loss_rv.detach(), | |
| 'discriminator_loss_gt': disc_loss_gt.detach(), | |
| 'discriminator_loss_grad': grad_penalty.detach() | |
| } | |
| def train(self): | |
| """Perform training.""" | |
| # archive code and configs | |
| if self.archive_code: | |
| misc.archive_code(osp.join(self.checkpoint_dir, 'archived_code.zip'), filetypes=['.py']) | |
| misc.dump_yaml(osp.join(self.checkpoint_dir, 'configs.yml'), self.cfgs) | |
| # initialize | |
| start_epoch = 0 | |
| self.total_iter = 0 | |
| self.total_iter = self.original_total_iter | |
| self.metrics_trace.reset() | |
| self.model.to(self.device) | |
| if self.model.enable_disc: | |
| self.model.reset_only_disc_optimizer() | |
| if self.few_shot_resume: | |
| resume_model_name = self.cfgs.get('few_shot_resume_name', None) | |
| start_epoch, self.total_iter = self.load_checkpoint(optim=True, checkpoint_name=resume_model_name) | |
| self.model.ddp(self.rank, self.world_size) | |
| # use tensorboard | |
| if self.use_logger: | |
| from torch.utils.tensorboard import SummaryWriter | |
| self.logger = SummaryWriter(osp.join(self.checkpoint_dir, 'logs', datetime.now().strftime("%Y%m%d-%H%M%S")), flush_secs=10) | |
| # self.viz_data_iterator = indefinite_generator_from_list(self.val_loader) if self.visualize_validation else indefinite_generator_from_list(self.train_loader) | |
| self.viz_data_iterator = indefinite_generator(self.val_loader[0]) if self.visualize_validation else indefinite_generator(self.train_loader[0]) | |
| if self.fix_viz_batch: | |
| self.viz_batch = next(self.viz_data_iterator) | |
| if self.test_loader is not None: | |
| self.viz_test_data_iterator = indefinite_generator(self.test_loader[0]) if self.visualize_validation else indefinite_generator(self.train_loader[0]) | |
| # run_epochs | |
| epoch = 0 | |
| for epoch in range(start_epoch, self.num_epochs): | |
| metrics = self.run_epoch(epoch) | |
| if self.combine_dataset: | |
| self.train_loader[0].dataset._shuffle_all() | |
| self.metrics_trace.append("train", metrics) | |
| if (epoch+1) % self.save_checkpoint_freq == 0: | |
| self.save_checkpoint(epoch+1, total_iter=self.total_iter, optim=True) | |
| # if self.cfgs.get('pyplot_metrics', True): | |
| # self.metrics_trace.plot(pdf_path=osp.join(self.checkpoint_dir, 'metrics.pdf')) | |
| self.metrics_trace.save(osp.join(self.checkpoint_dir, 'metrics.json')) | |
| print(f"Training completed for all {epoch+1} epochs.") | |
| def run_epoch(self, epoch): | |
| """Run one training epoch.""" | |
| metrics = self.make_metrics() | |
| self.model.set_train() | |
| max_loader_len = max([len(loader) for loader in self.train_loader]) | |
| train_generators = [indefinite_generator(loader) for loader in self.train_loader] | |
| iteration = 0 | |
| while iteration < max_loader_len * len(self.train_loader): | |
| for generator in train_generators: | |
| batch = next(generator) | |
| self.total_iter += 1 | |
| num_seqs, num_frames = batch[0].shape[:2] | |
| total_im_num = num_seqs * num_frames | |
| if self.enable_memory_bank: | |
| batch_features = self.forward_fix_embeddings(batch) | |
| batch_embedding, embeddings, weights = self.retrieve_memory_bank(batch_features, batch) | |
| bank_embedding_model_input = [batch_embedding, embeddings, weights] | |
| else: | |
| # bank_embedding_model_input = None | |
| batch_features = self.forward_fix_embeddings(batch) | |
| weights = { | |
| "weights": torch.rand(1,10).to(batch_features.device), | |
| "pick_idx": torch.randint(low=0, high=60, size=(1, 10)).to(batch_features.device) | |
| } | |
| bank_embedding_model_input = [batch_features[0], batch_features, weights] | |
| m = self.model.forward(batch, epoch=epoch, iter=iteration, total_iter=self.total_iter, which_data=self.dataset, is_training=True, bank_embedding=bank_embedding_model_input) | |
| # self.model.backward() | |
| self.optimizerFewShot.zero_grad() | |
| self.model.total_loss.backward() | |
| self.optimizerFewShot.step() | |
| self.model.total_loss = 0. | |
| # if self.cfgs.get('texture_way', None) is not None and self.cfgs.get('gan_tex', False): | |
| if self.model.few_shot_gan_tex: | |
| # the discriminator for local texture | |
| disc_ret = self.discriminator_texture_step() | |
| m.update(disc_ret) | |
| if self.model.enable_disc and (self.model.mask_discriminator_iter[0] < self.total_iter) and (self.model.mask_discriminator_iter[1] > self.total_iter): | |
| # the discriminator training | |
| discriminator_loss_dict, grad_loss = self.model.discriminator_step() | |
| m.update( | |
| { | |
| 'mask_disc_loss_discriminator': discriminator_loss_dict['discriminator_loss'] - grad_loss, | |
| 'mask_disc_loss_discriminator_grad': grad_loss, | |
| 'mask_disc_loss_discriminator_rv': discriminator_loss_dict['discriminator_loss_rv'], | |
| 'mask_disc_loss_discriminator_iv': discriminator_loss_dict['discriminator_loss_iv'], | |
| 'mask_disc_loss_discriminator_gt': discriminator_loss_dict['discriminator_loss_gt'] | |
| } | |
| ) | |
| self.logger.add_histogram('train_'+'discriminator_logits/random_view', discriminator_loss_dict['d_rv'], self.total_iter) | |
| if discriminator_loss_dict['d_iv'] is not None: | |
| self.logger.add_histogram('train_'+'discriminator_logits/input_view', discriminator_loss_dict['d_iv'], self.total_iter) | |
| if discriminator_loss_dict['d_gt'] is not None: | |
| self.logger.add_histogram('train_'+'discriminator_logits/gt_view', discriminator_loss_dict['d_gt'], self.total_iter) | |
| metrics.update(m, total_im_num) | |
| if self.rank == 0: | |
| print(f"T{epoch:04}/{iteration:05}/{metrics}") | |
| if self.iteration_save and self.total_iter % self.iteration_save_freq == 0: | |
| self.save_checkpoint(epoch+1, total_iter=self.total_iter, optim=True, use_iter=True) | |
| # ## reset optimizers | |
| # if self.cfgs.get('opt_reset_every_iter', 0) > 0 and self.total_iter < self.cfgs.get('opt_reset_end_iter', 0): | |
| # if self.total_iter % self.cfgs.get('opt_reset_every_iter', 0) == 0: | |
| # self.model.reset_optimizers() | |
| if misc.is_main_process() and self.use_logger: | |
| if self.rank == 0 and self.total_iter % self.log_freq_losses == 0: | |
| for name, loss in m.items(): | |
| label = f'cub_loss_train/{name[4:]}' if 'cub' in name else f'loss_train/{name}' | |
| self.logger.add_scalar(label, loss, self.total_iter) | |
| if self.rank == 0 and self.save_result_freq is not None and self.total_iter % self.save_result_freq == 0: | |
| with torch.no_grad(): | |
| m = self.model.forward(batch, epoch=epoch, iter=iteration, total_iter=self.total_iter, save_results=False, save_dir=self.train_result_dir, which_data=self.dataset, is_training=False, bank_embedding=bank_embedding_model_input) | |
| torch.cuda.empty_cache() | |
| if self.total_iter % self.log_freq_images == 0: | |
| with torch.no_grad(): | |
| if self.rank == 0 and self.log_train_images: | |
| m = self.model.forward(batch, epoch=epoch, iter=iteration, viz_logger=self.logger, total_iter=self.total_iter, which_data=self.dataset, logger_prefix='train_', is_training=False, bank_embedding=bank_embedding_model_input) | |
| if self.fix_viz_batch: | |
| print(f'fix_viz_batch:{self.fix_viz_batch}') | |
| batch_val = self.viz_batch | |
| else: | |
| batch_val = next(self.viz_data_iterator) | |
| if self.visualize_validation: | |
| import time | |
| vis_start = time.time() | |
| # batch = next(self.viz_data_iterator) | |
| # try: | |
| # batch = next(self.viz_data_iterator) | |
| # except: # iterator exhausted | |
| # self.reset_viz_data_iterator() | |
| # batch = next(self.viz_data_iterator) | |
| if self.enable_memory_bank: | |
| batch_features_val = self.forward_fix_embeddings(batch_val) | |
| batch_embedding_val, embeddings_val, weights_val = self.retrieve_memory_bank(batch_features_val, batch_val) | |
| bank_embedding_model_input_val = [batch_embedding_val, embeddings_val, weights_val] | |
| else: | |
| # bank_embedding_model_input_val = None | |
| batch_features_val = self.forward_fix_embeddings(batch_val) | |
| weights_val = { | |
| "weights": torch.rand(1,10).to(batch_features_val.device), | |
| "pick_idx": torch.randint(low=0, high=60, size=(1, 10)).to(batch_features_val.device) | |
| } | |
| bank_embedding_model_input_val = [batch_features_val[0], batch_features_val, weights_val] | |
| if self.total_iter % self.save_result_freq == 0: | |
| m = self.model.forward(batch_val, epoch=epoch, iter=iteration, viz_logger=self.logger, total_iter=self.total_iter, save_results=False, save_dir=self.train_result_dir, which_data=self.dataset, logger_prefix='val_', is_training=False, bank_embedding=bank_embedding_model_input_val) | |
| torch.cuda.empty_cache() | |
| vis_end = time.time() | |
| print(f"vis time: {vis_end - vis_start}") | |
| if self.test_loader is not None: | |
| # unseen category test visualization | |
| batch_test = next(self.viz_test_data_iterator) | |
| if self.enable_memory_bank: | |
| batch_features_test = self.forward_fix_embeddings(batch_test) | |
| batch_embedding_test, embeddings_test, weights_test = self.retrieve_memory_bank(batch_features_test, batch_test) | |
| bank_embedding_model_input_test = [batch_embedding_test, embeddings_test, weights_test] | |
| else: | |
| # bank_embedding_model_input_test = None | |
| batch_features_test = self.forward_fix_embeddings(batch_test) | |
| weights_test = { | |
| "weights": torch.rand(1,10).to(batch_features_test.device), | |
| "pick_idx": torch.randint(low=0, high=60, size=(1, 10)).to(batch_features_test.device) | |
| } | |
| bank_embedding_model_input_test = [batch_features_test[0], batch_features_test, weights_test] | |
| m_test = self.model.forward(batch_test, epoch=epoch, iter=iteration, viz_logger=self.logger, total_iter=self.total_iter, which_data=self.dataset, logger_prefix='test_', is_training=False, bank_embedding=bank_embedding_model_input_test) | |
| vis_test_end = time.time() | |
| print(f"vis test time: {vis_test_end - vis_end}") | |
| for name, loss in m_test.items(): | |
| if self.rank == 0: | |
| self.logger.add_scalar(f'loss_test/{name}', loss, self.total_iter) | |
| for name, loss in m.items(): | |
| if self.rank == 0: | |
| self.logger.add_scalar(f'loss_val/{name}', loss, self.total_iter) | |
| torch.cuda.empty_cache() | |
| iteration += 1 | |
| self.model.scheduler_step() | |
| return metrics | |