| |
| |
| |
| |
|
|
| import logging |
| import os |
| import random |
| import subprocess |
| from typing import Callable, List, Optional, Tuple |
|
|
| import numpy as np |
| import torch |
| from torch import Tensor, nn |
|
|
| logger = logging.getLogger("dinov3") |
|
|
|
|
| def cat_keep_shapes(x_list: List[Tensor]) -> Tuple[Tensor, List[Tuple[int]], List[int]]: |
| shapes = [x.shape for x in x_list] |
| num_tokens = [x.select(dim=-1, index=0).numel() for x in x_list] |
| flattened = torch.cat([x.flatten(0, -2) for x in x_list]) |
| return flattened, shapes, num_tokens |
|
|
|
|
| def uncat_with_shapes(flattened: Tensor, shapes: List[Tuple[int]], num_tokens: List[int]) -> List[Tensor]: |
| outputs_splitted = torch.split_with_sizes(flattened, num_tokens, dim=0) |
| shapes_adjusted = [shape[:-1] + torch.Size([flattened.shape[-1]]) for shape in shapes] |
| outputs_reshaped = [o.reshape(shape) for o, shape in zip(outputs_splitted, shapes_adjusted)] |
| return outputs_reshaped |
|
|
|
|
| def named_replace( |
| fn: Callable, |
| module: nn.Module, |
| name: str = "", |
| depth_first: bool = True, |
| include_root: bool = False, |
| ) -> nn.Module: |
| if not depth_first and include_root: |
| module = fn(module=module, name=name) |
| for child_name_o, child_module in list(module.named_children()): |
| child_name = ".".join((name, child_name_o)) if name else child_name_o |
| new_child = named_replace( |
| fn=fn, |
| module=child_module, |
| name=child_name, |
| depth_first=depth_first, |
| include_root=True, |
| ) |
| setattr(module, child_name_o, new_child) |
|
|
| if depth_first and include_root: |
| module = fn(module=module, name=name) |
| return module |
|
|
|
|
| def named_apply( |
| fn: Callable, |
| module: nn.Module, |
| name: str = "", |
| depth_first: bool = True, |
| include_root: bool = False, |
| ) -> nn.Module: |
| if not depth_first and include_root: |
| fn(module=module, name=name) |
| for child_name, child_module in module.named_children(): |
| child_name = ".".join((name, child_name)) if name else child_name |
| named_apply( |
| fn=fn, |
| module=child_module, |
| name=child_name, |
| depth_first=depth_first, |
| include_root=True, |
| ) |
| if depth_first and include_root: |
| fn(module=module, name=name) |
| return module |
|
|
|
|
| def fix_random_seeds(seed: int = 31): |
| """ |
| Fix random seeds. |
| """ |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| np.random.seed(seed) |
| random.seed(seed) |
|
|
|
|
| def get_sha() -> str: |
| cwd = os.path.dirname(os.path.abspath(__file__)) |
|
|
| def _run(command): |
| return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() |
|
|
| sha = "N/A" |
| diff = "clean" |
| branch = "N/A" |
| try: |
| sha = _run(["git", "rev-parse", "HEAD"]) |
| subprocess.check_output(["git", "diff"], cwd=cwd) |
| diff = _run(["git", "diff-index", "HEAD"]) |
| diff = "has uncommited changes" if diff else "clean" |
| branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) |
| except Exception: |
| pass |
| message = f"sha: {sha}, status: {diff}, branch: {branch}" |
| return message |
|
|
|
|
| def get_conda_env() -> Tuple[Optional[str], Optional[str]]: |
| conda_env_name = os.environ.get("CONDA_DEFAULT_ENV") |
| conda_env_path = os.environ.get("CONDA_PREFIX") |
| return conda_env_name, conda_env_path |
|
|
|
|
| def count_parameters(module: nn.Module) -> int: |
| c = 0 |
| for m in module.parameters(): |
| c += m.nelement() |
| return c |
|
|
|
|
| def has_batchnorms(model: nn.Module) -> bool: |
| bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm) |
| for _, module in model.named_modules(): |
| if isinstance(module, bn_types): |
| return True |
| return False |
|
|
