import datetime import io import os import random import sys import time from collections import defaultdict, deque from pathlib import Path import matplotlib import numpy as np import torch import torch.distributed as dist import torch.nn.functional as F from einops import rearrange from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD from timm.utils import get_state_dict from torch import inf # sys.path.append(os.path.join(os.environ['HOME'], '.cache/torch/CutLER')) # sys.path.append(os.path.join(os.environ['HOME'], '.cache/torch/CutLER/maskcut')) # sys.path.append(os.path.join(os.environ['HOME'], '.cache/torch/CutLER/third_party')) # import dino # from maskcut import get_affinity_matrix, second_smallest_eigenvector, get_salient_areas, check_num_fg_corners, get_masked_affinity_matrix # # #from maskcut import get_affinity_matrix, second_smallest_eigenvector, get_salient_areas, check_num_fg_corners # # DINO hyperparameters # global dino_backbone # dino_backbone = None def patchify(x, tubelet_size, patch_size): ''' :param x: [B, C, T, H, W] :param tubelet_size: 2 :param patch_size: (8, 8) :return: ''' videos_squeeze = rearrange(x, 'b c (t p0) (h p1) (w p2) -> b (t h w) (p0 p1 p2) c', p0=tubelet_size, p1=patch_size[0], p2=patch_size[1]) videos_patch = rearrange(videos_squeeze, 'b n p c -> b n (p c)') return videos_patch def imagenet_unnormalize(x, temporal_dim=2): device = x.device if len(x.shape) == 3: if x.shape[0] == 3: # "channel_first" mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[:, None, None].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[:, None, None].to(x) else: # channel_last mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[None, None, :].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[None, None, :].to(x) elif len(x.shape) == 4: mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[None, :, None, None].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[None, :, None, None].to(x) elif len(x.shape) == 5: mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[None, None, :, None, None].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[None, None, :, None, None].to(x) if temporal_dim == 2: mean = mean.transpose(1,2) std = std.transpose(1,2) return x * std + mean def imagenet_normalize(x, temporal_dim=2): device = x.device if len(x.shape) == 3: if x.shape[0] == 3: # "channel_first" mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[:, None, None].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[:, None, None].to(x) else: # channel_last mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[None, None, :].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[None, None, :].to(x) elif len(x.shape) == 4: mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[None, :, None, None].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[None, :, None, None].to(x) elif len(x.shape) == 5: mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[None, None, :, None, None].to(x) std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[None, None, :, None, None].to(x) if temporal_dim == 2: mean = mean.transpose(1,2) std = std.transpose(1,2) return (x - mean) / std def sinusoidal_embedding(x, n_freq=5, keep_ori=True): """ create sin embedding for 3d vectors input: x: *x3 n_freq: number of raised frequency """ shape = list(x.shape) assert x.shape[-1] == 3, "expect the last dimension to have size 3" x = x.reshape(-1, 3) embedded = [] if keep_ori: embedded.append(x) emb_fns = [torch.sin, torch.cos] freqs = 2. ** torch.linspace(0., n_freq - 1, steps=n_freq) for freq in freqs: for emb_fn in emb_fns: embedded.append(emb_fn(freq * x)) embedded = torch.cat(embedded, dim=-1) C = embedded.shape[-1] embedded = embedded.reshape(shape[:-1] + [C]) return embedded class SmoothedValue(object): """Track a series of values and provide access to smoothed values over a window or the global series average. """ def __init__(self, window_size=20, fmt=None): if fmt is None: fmt = "{median:.4f} ({global_avg:.4f})" self.deque = deque(maxlen=window_size) self.total = 0.0 self.count = 0 self.fmt = fmt def update(self, value, n=1): self.deque.append(value) self.count += n self.total += value * n def synchronize_between_processes(self): """ Warning: does not synchronize the deque! """ if not is_dist_avail_and_initialized(): return t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda') dist.barrier() dist.all_reduce(t) t = t.tolist() self.count = int(t[0]) self.total = t[1] @property def median(self): d = torch.tensor(list(self.deque)) return d.median().item() @property def avg(self): d = torch.tensor(list(self.deque), dtype=torch.float32) return d.mean().item() @property def global_avg(self): return self.total / self.count @property def max(self): return max(self.deque) @property def value(self): return self.deque[-1] def __str__(self): return self.fmt.format( median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value) class MetricLogger(object): def __init__(self, delimiter="\t"): self.meters = defaultdict(SmoothedValue) self.delimiter = delimiter def update(self, **kwargs): for k, v in kwargs.items(): if v is None: continue if isinstance(v, torch.Tensor): v = v.item() assert isinstance(v, (float, int)) self.meters[k].update(v) def update2(self, kwargs): for k, v in kwargs.items(): if v is None: continue if isinstance(v, torch.Tensor): v = v.item() assert isinstance(v, (float, int)) self.meters[k].update(v) def __getattr__(self, attr): if attr in self.meters: return self.meters[attr] if attr in self.__dict__: return self.__dict__[attr] raise AttributeError("'{}' object has no attribute '{}'".format( type(self).__name__, attr)) def __str__(self): loss_str = [] for name, meter in self.meters.items(): loss_str.append( "{}: {}".format(name, str(meter)) ) return self.delimiter.join(loss_str) def synchronize_between_processes(self): for meter in self.meters.values(): meter.synchronize_between_processes() def add_meter(self, name, meter): self.meters[name] = meter def log_every(self, iterable, print_freq, header=None): i = 0 if not header: header = '' start_time = time.time() end = time.time() iter_time = SmoothedValue(fmt='{avg:.2f}') data_time = SmoothedValue(fmt='{avg:.4f}') space_fmt = ':' + str(len(str(len(iterable)))) + 'd' log_msg = [ header, '[{0' + space_fmt + '}/{1}]', 'eta: {eta}', '{meters}', 'time: {time}', 'data: {data}' ] if torch.cuda.is_available(): log_msg.append('max mem: {memory:.0f}') log_msg = self.delimiter.join(log_msg) MB = 1024.0 * 1024.0 for obj in iterable: data_time.update(time.time() - end) yield obj iter_time.update(time.time() - end) if i % print_freq == 0 or i == len(iterable) - 1: eta_seconds = iter_time.global_avg * (len(iterable) - i) eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) if torch.cuda.is_available(): print(log_msg.format( i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time), memory=torch.cuda.max_memory_allocated() / MB)) else: print(log_msg.format( i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time))) i += 1 end = time.time() total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('{} Total time: {} ({:.6f} s / it)'.format( header, total_time_str, total_time / len(iterable))) def seed_worker(worker_id): worker_seed = torch.initial_seed() % 2**32 np.random.seed(worker_seed) random.seed(worker_seed) def _load_checkpoint_for_ema(model_ema, checkpoint): """ Workaround for ModelEma._load_checkpoint to accept an already-loaded object """ mem_file = io.BytesIO() torch.save(checkpoint, mem_file) mem_file.seek(0) model_ema._load_checkpoint(mem_file) def setup_for_distributed(is_master): """ This function disables printing when not in master process """ import builtins as __builtin__ builtin_print = __builtin__.print def print(*args, **kwargs): force = kwargs.pop('force', False) if is_master or force: builtin_print(*args, **kwargs) __builtin__.print = print def is_dist_avail_and_initialized(): if not dist.is_available(): return False if not dist.is_initialized(): return False return True def get_world_size(): if not is_dist_avail_and_initialized(): return 1 return dist.get_world_size() def get_rank(): if not is_dist_avail_and_initialized(): return 0 return dist.get_rank() def is_main_process(): return get_rank() == 0 def save_on_master(*args, **kwargs): if is_main_process(): torch.save(*args, **kwargs) def init_distributed_mode(args): args.distributed = True args.rank = int(os.environ["RANK"]) args.gpu = int(os.environ['LOCAL_RANK']) args.world_size = int(os.environ['WORLD_SIZE']) args.dist_backend = 'nccl' torch.distributed.init_process_group( backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank ) torch.distributed.barrier() setup_for_distributed(args.rank == 0) def load_state_dict(model, state_dict, prefix='', ignore_missing="relative_position_index"): missing_keys = [] unexpected_keys = [] error_msgs = [] metadata = getattr(state_dict, '_metadata', None) state_dict = state_dict.copy() if metadata is not None: state_dict._metadata = metadata def load(module, prefix=''): local_metadata = {} if metadata is None else metadata.get( prefix[:-1], {}) module._load_from_state_dict( state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs) for name, child in module._modules.items(): if child is not None: load(child, prefix + name + '.') load(model, prefix=prefix) warn_missing_keys = [] ignore_missing_keys = [] for key in missing_keys: keep_flag = True for ignore_key in ignore_missing.split('|'): if ignore_key in key: keep_flag = False break if keep_flag: warn_missing_keys.append(key) else: ignore_missing_keys.append(key) missing_keys = warn_missing_keys if len(missing_keys) > 0: print("Weights of {} not initialized from pretrained model: {}".format( model.__class__.__name__, missing_keys)) if len(unexpected_keys) > 0: print("Weights from pretrained model not used in {}: {}".format( model.__class__.__name__, unexpected_keys)) if len(ignore_missing_keys) > 0: print("Ignored weights of {} not initialized from pretrained model: {}".format( model.__class__.__name__, ignore_missing_keys)) if len(error_msgs) > 0: print('\n'.join(error_msgs)) class NativeScalerWithGradNormCount: state_dict_key = "amp_scaler" def __init__(self): self._scaler = torch.cuda.amp.GradScaler() def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True): self._scaler.scale(loss).backward(create_graph=create_graph) # breakpoint() if update_grad: if clip_grad is not None: assert parameters is not None self._scaler.unscale_(optimizer) # unscale the gradients of optimizer's assigned params in-place norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad) else: self._scaler.unscale_(optimizer) norm = get_grad_norm_(parameters) self._scaler.step(optimizer) self._scaler.update() else: norm = None return norm def state_dict(self): return self._scaler.state_dict() def load_state_dict(self, state_dict): self._scaler.load_state_dict(state_dict) def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor: if isinstance(parameters, torch.Tensor): parameters = [parameters] parameters = [p for p in parameters if p.grad is not None] norm_type = float(norm_type) if len(parameters) == 0: return torch.tensor(0.) device = parameters[0].grad.device if norm_type == inf: total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters) else: total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type) return total_norm def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, start_warmup_value=0, warmup_steps=-1): warmup_schedule = np.array([]) warmup_iters = warmup_epochs * niter_per_ep if warmup_steps > 0: warmup_iters = warmup_steps if warmup_epochs > 0: warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters) iters = np.arange(epochs * niter_per_ep - warmup_iters) iter_per_len = iters/len(iters) schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iter_per_len)) # schedule = np.array( # [final_value + 0.5 * (base_value - final_value) * (1 + math.cos(math.pi * i / (len(iters)))) for i in iters]) schedule = np.concatenate((warmup_schedule, schedule)) assert len(schedule) == epochs * niter_per_ep return schedule def get_model_num_parameters(model): num_parameters = sum([v.numel() for v in model.parameters() if v.requires_grad]) human_readable_fn = lambda num: \ f'{num / 1e9:.3f} B' if num >= 1e9 else f'{num / 1e6:.3f} M' \ if num >= 1e6 else f'{num / 1e3:.3f} K' if num >= 1e3 else str(num) num_parameters_str = human_readable_fn(num_parameters) return num_parameters, num_parameters_str def save_model(args, epoch, model, optimizer, loss_scaler, model_ema=None): output_dir = Path(args.output_dir) epoch_name = str(epoch) if loss_scaler is not None: checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)] for checkpoint_path in checkpoint_paths: to_save = { 'model': model.module.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'scaler': loss_scaler.state_dict(), 'args': args, } if model_ema is not None: to_save['model_ema'] = get_state_dict(model_ema) save_on_master(to_save, checkpoint_path) else: client_state = {'epoch': epoch} if model_ema is not None: client_state['model_ema'] = get_state_dict(model_ema) model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state) def auto_load_model(args, model, optimizer, loss_scaler, model_ema=None, global_rank=None): output_dir = Path(args.output_dir) if loss_scaler is not None: # torch.amp if len(args.resume) == 0: import glob if global_rank is None: all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*.pth')) else: all_checkpoints = glob.glob(os.path.join(output_dir, f'checkpoint-*-rank-{global_rank}.pth')) latest_ckpt = -1 for ckpt in all_checkpoints: if global_rank is None: t = ckpt.split('-')[-1].split('.')[0] else: t = ckpt.split('checkpoint-')[1].split('-')[0] if t.isdigit(): latest_ckpt = max(int(t), latest_ckpt) if latest_ckpt >= 0: if global_rank is None: args.resume = os.path.join(output_dir, 'checkpoint-%d.pth' % latest_ckpt) else: args.resume = os.path.join(output_dir, 'checkpoint-%d-rank-%d.pth' % (latest_ckpt, global_rank)) if args.resume: print("Auto resume checkpoint: %s" % args.resume) if args.resume: if args.resume.startswith('https'): checkpoint = torch.hub.load_state_dict_from_url( args.resume, map_location='cpu', check_hash=True) else: checkpoint = torch.load(args.resume, map_location='cpu') model.module.load_state_dict(checkpoint['model']) print("Resume checkpoint %s" % args.resume) if 'optimizer' in checkpoint and 'epoch' in checkpoint: optimizer.load_state_dict(checkpoint['optimizer']) args.start_epoch = checkpoint['epoch'] + 1 if hasattr(args, 'model_ema') and args.model_ema: _load_checkpoint_for_ema(model_ema, checkpoint['model_ema']) if 'scaler' in checkpoint: loss_scaler.load_state_dict(checkpoint['scaler']) else: # deepspeed, only support '--auto_resume'. import glob all_checkpoints = glob.glob(os.path.join(output_dir, 'checkpoint-*')) latest_ckpt = -1 for ckpt in all_checkpoints: t = ckpt.split('-')[-1].split('.')[0] if t.isdigit(): latest_ckpt = max(int(t), latest_ckpt) if latest_ckpt >= 0: args.resume = os.path.join(output_dir, 'checkpoint-%d' % latest_ckpt) print("Auto resume checkpoint: %d" % latest_ckpt) _, client_states = model.load_checkpoint(args.output_dir, tag='checkpoint-%d' % latest_ckpt) args.start_epoch = client_states['epoch'] + 1 if model_ema is not None: if args.model_ema: _load_checkpoint_for_ema(model_ema, client_states['model_ema']) def unpatchify(x, patch_size): """ x: (N, L, patch_size**2*3) imgs: (N, 3, H, W) """ p = patch_size h = w = int(x.shape[1] ** .5) assert h * w == x.shape[1] x = x.reshape(shape=(x.shape[0], h, w, p, p, 3)) x = torch.einsum('nhwpqc->nchpwq', x) imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p)) return imgs def unpatchify_cwm(x, patch_size, mask=None): """ x: (N, L, patch_size**2 *3) imgs: (N, 3, H, W) """ if mask is not None: h = w = int(mask.shape[1] ** .5) recon = torch.zeros(x.shape[0], h*w, x.shape[-1]).to(x) recon[mask] = x.flatten(0, 1) else: h = w = int(x.shape[1] ** .5) recon = x p = patch_size assert h * w == recon.shape[1] recon = recon.reshape(shape=(recon.shape[0], h, w, p, p, 3)) recon = torch.einsum('nhwpqc->nchpwq', recon) imgs = recon.reshape(shape=(recon.shape[0], 3, h * p, h * p)) return imgs def sample_embedding(embedding, pos, mode='bilinear'): """ Sample embedding tensor at specified positions embedding: [B, H, W, C] pos: [B, P, 2] (convention: first dim is row, second dim is column) """ embedding = embedding.permute(0, 3, 1, 2) # [B, C, H, W] device = embedding.device # grid_sampling assues first value to be column-dimension, second value to be row-dimension pos = pos.flip(dims=(-1,)) assert pos.min() >= -1 and pos.max() <= 1, "grid sampling expect to be in range [-1, 1]" return F.grid_sample(embedding, pos[:, None].to(device), mode=mode).squeeze(-2).permute(0, 2, 1) # [B, P, C] def sample_positions_from_dist(size, dist): """ Samples positions from a given unnormalized probability distribution. Parameters: num (int): The number of samples to draw for each distribution in the batch. dist (torch.Tensor): A float tensor of shape [B, H, W] representing the unnormalized probability distributions for B batches each of length N. Returns: torch.Tensor: A tensor of shape [B, num] containing the sampled positions. """ assert dist.dim() == 3, "dist should be a 3D tensor with shape [B, H, W]." assert len(size) == 2, "size should be a 2D tuple (batch_size, num_samples)" B, H, W = dist.shape new_B, num_samples = size if dist.min() < 0: dist -= dist.min() # Flatten the last two dimensions to make it [B, H*W] flattened_dist = dist.view(B, -1) # Sample indices according to the normalized distribution sampled_indices = torch.multinomial(flattened_dist, new_B * num_samples, replacement=True) # Convert the flattened indices back to 2D indices sampled_row_indices = sampled_indices // W sampled_col_indices = sampled_indices % W # Stack the row and column indices samples = torch.stack((sampled_row_indices, sampled_col_indices), dim=-1) samples = samples.view(new_B, num_samples, 2) return samples # # def get_dino_predominance(images, dims=[28, 28], current_mask=None, painting=None, img_size=[224, 224]): # global dino_backbone # if dino_backbone is None: # vit_arch = 'base' # vit_feat = 'k' # patch_size = 8 # # DINO pre-trained model # url = "https://dl.fbaipublicfiles.com/dino/dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth" # feat_dim = 768 # dino_backbone = dino.ViTFeat(url, feat_dim, vit_arch, vit_feat, patch_size) # dino_backbone = dino_backbone.eval().requires_grad_(False).cuda() # # input_dino = images # # input_dino = input_dino - torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(input_dino.device) # # input_dino = input_dino / torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(input_dino.device) # # input_dino = images.tensor # input_dino = torch.nn.functional.interpolate(input_dino, size=img_size, mode='bilinear') # features = dino_backbone(input_dino) # # predominence_map = [] # # for i in range(features.shape[0]): # feats = features[i] # if current_mask == None: # painting = torch.from_numpy(np.zeros(dims)) # painting = painting.to(feats) # else: # feats, painting = get_masked_affinity_matrix(painting, feats, current_mask, ps=dims[0]) # # A, D = get_affinity_matrix(feats, tau=0.15) # # get the second-smallest eigenvector # _, second_smallest_vec = second_smallest_eigenvector(A, D) # # get salient area # bipartition = get_salient_areas(second_smallest_vec) # # # check if we should reverse the partition based on: # # 1) peak of the 2nd smallest eigvec 2) object centric bias # seed = np.argmax(np.abs(second_smallest_vec)) # nc = check_num_fg_corners(bipartition, dims) # if nc >= 2: # reverse = True # else: # reverse = bipartition[seed] != 1 # if reverse: # second_smallest_vec = 1 - second_smallest_vec # second_smallest_vec = torch.tensor(second_smallest_vec).to(images.device).contiguous() # map = torch.nn.functional.interpolate(second_smallest_vec.reshape(1, 1, dims[0], dims[1]), size=img_size, # mode='bilinear') # map -= map.min() # map /= map.max() # predominence_map.append(map) # init_dist = torch.cat(predominence_map, dim=0).detach() # return init_dist, A, feats, painting def interpolate_pos_encoding(pos_embed, n_frames, h, w): N = pos_embed.shape[1] if N == (h * w * n_frames): return pos_embed old_h = old_w = int((N / n_frames) ** 0.5) patch_pos_embed = pos_embed.view(1, n_frames, old_h, old_w, -1).flatten(0, 1).permute(0, 3, 1, 2) patch_pos_embed = F.interpolate( patch_pos_embed, size=(h, w), mode='bicubic', ) return patch_pos_embed.permute(0, 2, 3, 1).flatten(0, 2).unsqueeze(0) def flow_to_rgb(vec, flow_mag_range=None, white_bg=False): height, width = vec.shape[:2] scaling = 50. / (height**2 + width**2)**0.5 direction = (np.arctan2(vec[..., 0], vec[..., 1]) + np.pi) / (2 * np.pi) norm = np.linalg.norm(vec, axis=-1) if flow_mag_range is None: flow_mag_range = norm.min(), norm.max() magnitude = np.clip((norm - flow_mag_range[0]) * scaling, 0., 1.) if white_bg == True: value = np.ones_like(direction) hsv = np.stack([direction, magnitude, saturation], axis=-1) else: saturation = np.ones_like(direction) hsv = np.stack([direction, saturation , magnitude], axis=-1) rgb = matplotlib.colors.hsv_to_rgb(hsv) return rgb