import os import math import torch import logging import subprocess import numpy as np import torch.distributed as dist # from torch._six import inf from torch import inf from PIL import Image from typing import Union, Iterable from collections import OrderedDict _tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]] ################################################################################# # Training Helper Functions # ################################################################################# ################################################################################# # Training Clip Gradients # ################################################################################# def get_grad_norm( parameters: _tensor_or_tensors, norm_type: float = 2.0) -> torch.Tensor: r""" Copy from torch.nn.utils.clip_grad_norm_ Clips gradient norm of an iterable of parameters. The norm is computed over all gradients together, as if they were concatenated into a single vector. Gradients are modified in-place. Args: parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a single Tensor that will have gradients normalized max_norm (float or int): max norm of the gradients norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for infinity norm. error_if_nonfinite (bool): if True, an error is thrown if the total norm of the gradients from :attr:`parameters` is ``nan``, ``inf``, or ``-inf``. Default: False (will switch to True in the future) Returns: Total norm of the parameter gradients (viewed as a single vector). """ if isinstance(parameters, torch.Tensor): parameters = [parameters] grads = [p.grad for p in parameters if p.grad is not None] norm_type = float(norm_type) if len(grads) == 0: return torch.tensor(0.) device = grads[0].device if norm_type == inf: norms = [g.detach().abs().max().to(device) for g in grads] total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms)) else: total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type) return total_norm def clip_grad_norm_( parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0, error_if_nonfinite: bool = False, clip_grad = True) -> torch.Tensor: r""" Copy from torch.nn.utils.clip_grad_norm_ Clips gradient norm of an iterable of parameters. The norm is computed over all gradients together, as if they were concatenated into a single vector. Gradients are modified in-place. Args: parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a single Tensor that will have gradients normalized max_norm (float or int): max norm of the gradients norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for infinity norm. error_if_nonfinite (bool): if True, an error is thrown if the total norm of the gradients from :attr:`parameters` is ``nan``, ``inf``, or ``-inf``. Default: False (will switch to True in the future) Returns: Total norm of the parameter gradients (viewed as a single vector). """ if isinstance(parameters, torch.Tensor): parameters = [parameters] grads = [p.grad for p in parameters if p.grad is not None] max_norm = float(max_norm) norm_type = float(norm_type) if len(grads) == 0: return torch.tensor(0.) device = grads[0].device if norm_type == inf: norms = [g.detach().abs().max().to(device) for g in grads] total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms)) else: total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type) # print(total_norm) if clip_grad: if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()): raise RuntimeError( f'The total norm of order {norm_type} for gradients from ' '`parameters` is non-finite, so it cannot be clipped. To disable ' 'this error and scale the gradients by the non-finite norm anyway, ' 'set `error_if_nonfinite=False`') clip_coef = max_norm / (total_norm + 1e-6) # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization # when the gradients do not reside in CPU memory. clip_coef_clamped = torch.clamp(clip_coef, max=1.0) for g in grads: g.detach().mul_(clip_coef_clamped.to(g.device)) # gradient_cliped = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type) # print(gradient_cliped) return total_norm ################################################################################# # Training Logger # ################################################################################# def create_logger(logging_dir): """ Create a logger that writes to a log file and stdout. """ if dist.get_rank() == 0: # real logger logging.basicConfig( level=logging.INFO, # format='[\033[34m%(asctime)s\033[0m] %(message)s', format='[%(asctime)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S', handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")] ) logger = logging.getLogger(__name__) else: # dummy logger (does nothing) logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) return logger def create_accelerate_logger(logging_dir, is_main_process=False): """ Create a logger that writes to a log file and stdout. """ if is_main_process: # real logger logging.basicConfig( level=logging.INFO, # format='[\033[34m%(asctime)s\033[0m] %(message)s', format='[%(asctime)s] %(message)s', datefmt='%Y-%m-%d %H:%M:%S', handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")] ) logger = logging.getLogger(__name__) else: # dummy logger (does nothing) logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) return logger def create_tensorboard(tensorboard_dir): """ Create a tensorboard that saves losses. """ if dist.get_rank() == 0: # real tensorboard # tensorboard writer = SummaryWriter(tensorboard_dir) return writer def write_tensorboard(writer, *args): ''' write the loss information to a tensorboard file. Only for pytorch DDP mode. ''' if dist.get_rank() == 0: # real tensorboard writer.add_scalar(args[0], args[1], args[2]) ################################################################################# # EMA Update/ DDP Training Utils # ################################################################################# @torch.no_grad() def update_ema(ema_model, model, decay=0.9999): """ Step the EMA model towards the current model. """ ema_params = OrderedDict(ema_model.named_parameters()) model_params = OrderedDict(model.named_parameters()) for name, param in model_params.items(): # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) def requires_grad(model, flag=True): """ Set requires_grad flag for all parameters in a model. """ for p in model.parameters(): p.requires_grad = flag def cleanup(): """ End DDP training. """ dist.destroy_process_group() def setup_distributed(backend="nccl", port=None): """Initialize distributed training environment. support both slurm and torch.distributed.launch see torch.distributed.init_process_group() for more details """ num_gpus = torch.cuda.device_count() print(f'Hahahahahaha') if "SLURM_JOB_ID" in os.environ: rank = int(os.environ["SLURM_PROCID"]) world_size = int(os.environ["SLURM_NTASKS"]) node_list = os.environ["SLURM_NODELIST"] addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1") # specify master port if port is not None: os.environ["MASTER_PORT"] = str(port) elif "MASTER_PORT" not in os.environ: # os.environ["MASTER_PORT"] = "29566" os.environ["MASTER_PORT"] = str(29566 + num_gpus) if "MASTER_ADDR" not in os.environ: os.environ["MASTER_ADDR"] = addr os.environ["WORLD_SIZE"] = str(world_size) os.environ["LOCAL_RANK"] = str(rank % num_gpus) os.environ["RANK"] = str(rank) else: rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) # torch.cuda.set_device(rank % num_gpus) print(f'before dist.init_process_group') dist.init_process_group( backend=backend, world_size=world_size, rank=rank, ) print(f'after dist.init_process_group') ################################################################################# # Testing Utils # ################################################################################# def save_video_grid(video, nrow=None): b, t, h, w, c = video.shape if nrow is None: nrow = math.ceil(math.sqrt(b)) ncol = math.ceil(b / nrow) padding = 1 video_grid = torch.zeros((t, (padding + h) * nrow + padding, (padding + w) * ncol + padding, c), dtype=torch.uint8) print(video_grid.shape) for i in range(b): r = i // ncol c = i % ncol start_r = (padding + h) * r start_c = (padding + w) * c video_grid[:, start_r:start_r + h, start_c:start_c + w] = video[i] return video_grid ################################################################################# # MMCV Utils # ################################################################################# def collect_env(): # Copyright (c) OpenMMLab. All rights reserved. from mmcv.utils import collect_env as collect_base_env from mmcv.utils import get_git_hash """Collect the information of the running environments.""" env_info = collect_base_env() env_info['MMClassification'] = get_git_hash()[:7] for name, val in env_info.items(): print(f'{name}: {val}') print(torch.cuda.get_arch_list()) print(torch.version.cuda) ################################################################################# # Long video generation Utils # ################################################################################# def mask_generation(mask_type, shape, dtype, device): b, c, f, h, w = shape if mask_type.startswith('random'): num = float(mask_type.split('random')[-1]) mask_f = torch.ones(1, 1, f, 1, 1, dtype=dtype, device=device) indices = torch.randperm(f, device=device)[:int(f*num)] mask_f[0, 0, indices, :, :] = 0 mask = mask_f.expand(b, c, -1, h, w) elif mask_type.startswith('first'): num = int(mask_type.split('first')[-1]) mask_f = torch.cat([torch.zeros(1, 1, num, 1, 1, dtype=dtype, device=device), torch.ones(1, 1, f-num, 1, 1, dtype=dtype, device=device)], dim=2) mask = mask_f.expand(b, c, -1, h, w) else: raise ValueError(f"Invalid mask type: {mask_type}") return mask def mask_generation_before(mask_type, shape, dtype, device): b, f, c, h, w = shape if mask_type.startswith('random'): num = float(mask_type.split('random')[-1]) mask_f = torch.ones(1, f, 1, 1, 1, dtype=dtype, device=device) indices = torch.randperm(f, device=device)[:int(f*num)] mask_f[0, indices, :, :, :] = 0 mask = mask_f.expand(b, -1, c, h, w) elif mask_type.startswith('first'): num = int(mask_type.split('first')[-1]) mask_f = torch.cat([torch.zeros(1, num, 1, 1, 1, dtype=dtype, device=device), torch.ones(1, f-num, 1, 1, 1, dtype=dtype, device=device)], dim=1) mask = mask_f.expand(b, -1, c, h, w) elif mask_type.startswith('uniform'): p = float(mask_type.split('uniform')[-1]) mask_f = torch.ones(1, f, 1, 1, 1, dtype=dtype, device=device) mask_f[0, torch.rand(f, device=device) < p, :, :, :] = 0 print(f'mask_f: = {mask_f}') mask = mask_f.expand(b, -1, c, h, w) print(f'mask.shape: = {mask.shape}, mask: = {mask}') elif mask_type.startswith('all'): mask = torch.ones(b,f,c,h,w,dtype=dtype,device=device) elif mask_type.startswith('onelast'): num = int(mask_type.split('onelast')[-1]) mask_one = torch.zeros(1,1,1,1,1, dtype=dtype, device=device) mask_mid = torch.ones(1,f-2*num,1,1,1,dtype=dtype, device=device) mask_last = torch.zeros_like(mask_one) mask = torch.cat([mask_one]*num + [mask_mid] + [mask_last]*num, dim=1) # breakpoint() mask = mask.expand(b, -1, c, h, w) elif mask_type.startswith('interpolate'): mask_f = [] for i in range(4): mask_zero = torch.zeros(1,1,1,1,1, dtype=dtype, device=device) mask_f.append(mask_zero) mask_one = torch.ones(1,3,1,1,1, dtype=dtype, device=device) mask_f.append(mask_one) mask = torch.cat(mask_f, dim=1) print(f'mask={mask}') elif mask_type.startswith('tsr'): mask_f = [] mask_zero = torch.zeros(1,1,1,1,1, dtype=dtype, device=device) mask_one = torch.ones(1,3,1,1,1, dtype=dtype, device=device) for i in range(15): mask_f.append(mask_zero) # not masked mask_f.append(mask_one) # masked mask_f.append(mask_zero) # not masked mask = torch.cat(mask_f, dim=1) # print(f'before mask.shape = {mask.shape}, mask = {mask}') # [1, 61, 1, 1, 1] mask = mask.expand(b, -1, c, h, w) # print(f'after mask.shape = {mask.shape}, mask = {mask}') # [4, 61, 3, 256, 256] else: raise ValueError(f"Invalid mask type: {mask_type}") return mask