Upload folder using huggingface_hub

utils/__init__.py

@@ -0,0 +1,7 @@
+from .criterion import *
+from .distributed import *
+from .init import *
+from .lr_scheduler import *
+from .metric import *
+from .misc import *
+from .profile import *

utils/criterion.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn.functional as F
+
+__all__ = ["label_smooth", "CrossEntropyWithSoftTarget", "CrossEntropyWithLabelSmooth"]
+
+
+def label_smooth(
+    target: torch.Tensor, n_classes: int, smooth_factor=0.1
+) -> torch.Tensor:
+    # convert to one-hot
+    batch_size = target.shape[0]
+    target = torch.unsqueeze(target, 1)
+    soft_target = torch.zeros((batch_size, n_classes), device=target.device)
+    soft_target.scatter_(1, target, 1)
+    # label smoothing
+    soft_target = torch.add(
+        soft_target * (1 - smooth_factor), smooth_factor / n_classes
+    )
+    return soft_target
+
+
+class CrossEntropyWithSoftTarget:
+    @staticmethod
+    def get_loss(pred: torch.Tensor, soft_target: torch.Tensor) -> torch.Tensor:
+        return torch.mean(
+            torch.sum(-soft_target * F.log_softmax(pred, dim=-1, _stacklevel=5), 1)
+        )
+
+    def __call__(self, pred: torch.Tensor, soft_target: torch.Tensor) -> torch.Tensor:
+        return self.get_loss(pred, soft_target)
+
+
+class CrossEntropyWithLabelSmooth:
+    def __init__(self, smooth_ratio=0.1):
+        super(CrossEntropyWithLabelSmooth, self).__init__()
+        self.smooth_ratio = smooth_ratio
+
+    def __call__(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        soft_target = label_smooth(target, pred.shape[1], self.smooth_ratio)
+        return CrossEntropyWithSoftTarget.get_loss(pred, soft_target)

utils/distributed.py

@@ -0,0 +1,55 @@
+from typing import List, Optional, Union
+
+import torch
+import torch.distributed
+from torchpack import distributed
+
+from utils.misc import list_mean, list_sum
+
+__all__ = ["ddp_reduce_tensor", "DistributedMetric"]
+
+
+def ddp_reduce_tensor(
+    tensor: torch.Tensor, reduce="mean"
+) -> Union[torch.Tensor, List[torch.Tensor]]:
+    tensor_list = [torch.empty_like(tensor) for _ in range(distributed.size())]
+    torch.distributed.all_gather(tensor_list, tensor.contiguous(), async_op=False)
+    if reduce == "mean":
+        return list_mean(tensor_list)
+    elif reduce == "sum":
+        return list_sum(tensor_list)
+    elif reduce == "cat":
+        return torch.cat(tensor_list, dim=0)
+    elif reduce == "root":
+        return tensor_list[0]
+    else:
+        return tensor_list
+
+
+class DistributedMetric(object):
+    """Average metrics for distributed training."""
+
+    def __init__(self, name: Optional[str] = None, backend="ddp"):
+        self.name = name
+        self.sum = 0
+        self.count = 0
+        self.backend = backend
+
+    def update(self, val: Union[torch.Tensor, int, float], delta_n=1):
+        val *= delta_n
+        if type(val) in [int, float]:
+            val = torch.Tensor(1).fill_(val).cuda()
+        if self.backend == "ddp":
+            self.count += ddp_reduce_tensor(
+                torch.Tensor(1).fill_(delta_n).cuda(), reduce="sum"
+            )
+            self.sum += ddp_reduce_tensor(val.detach(), reduce="sum")
+        else:
+            raise NotImplementedError
+
+    @property
+    def avg(self):
+        if self.count == 0:
+            return torch.Tensor(1).fill_(-1)
+        else:
+            return self.sum / self.count

utils/init.py

@@ -0,0 +1,77 @@
+import math
+from typing import Dict, List, Union
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.batchnorm import _BatchNorm
+
+__all__ = ["init_modules", "load_state_dict"]
+
+
+def init_modules(
+    module: Union[nn.Module, List[nn.Module]], init_type="he_fout"
+) -> None:
+    init_params = init_type.split("@")
+    if len(init_params) > 1:
+        init_params = float(init_params[1])
+    else:
+        init_params = None
+
+    if isinstance(module, list):
+        for sub_module in module:
+            init_modules(sub_module)
+    else:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                if init_type == "he_fout":
+                    n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                    m.weight.data.normal_(0, math.sqrt(2.0 / n))
+                elif init_type.startswith("kaiming_uniform"):
+                    nn.init.kaiming_uniform_(m.weight, a=math.sqrt(init_params or 5))
+                else:
+                    nn.init.kaiming_uniform_(m.weight, a=math.sqrt(init_params or 5))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, _BatchNorm):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            else:
+                weight = getattr(m, "weight", None)
+                bias = getattr(m, "bias", None)
+                if isinstance(weight, torch.nn.Parameter):
+                    nn.init.kaiming_uniform_(m.weight, a=math.sqrt(init_params or 5))
+                if isinstance(bias, torch.nn.Parameter):
+                    bias.data.zero_()
+
+
+def load_state_dict(
+    model: nn.Module, state_dict: Dict[str, torch.Tensor], strict=True
+) -> None:
+    current_state_dict = model.state_dict()
+    for key in state_dict:
+        if current_state_dict[key].shape != state_dict[key].shape:
+            if strict:
+                raise ValueError(
+                    "%s shape mismatch (src=%s, target=%s)"
+                    % (
+                        key,
+                        list(state_dict[key].shape),
+                        list(current_state_dict[key].shape),
+                    )
+                )
+            else:
+                print(
+                    "Skip loading %s due to shape mismatch (src=%s, target=%s)"
+                    % (
+                        key,
+                        list(state_dict[key].shape),
+                        list(current_state_dict[key].shape),
+                    )
+                )
+        else:
+            current_state_dict[key].copy_(state_dict[key])
+    model.load_state_dict(current_state_dict)

utils/lr_scheduler.py

@@ -0,0 +1,38 @@
+import math
+from typing import List
+
+import torch
+from torch.optim import Optimizer
+
+__all__ = ["CosineLRwithWarmup"]
+
+
+class CosineLRwithWarmup(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        warmup_steps: int,
+        warmup_lr: float,
+        decay_steps: int,
+        last_epoch: int = -1,
+    ) -> None:
+        self.warmup_steps = warmup_steps
+        self.warmup_lr = warmup_lr
+        self.decay_steps = decay_steps
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self) -> List[float]:
+        if self.last_epoch < self.warmup_steps:
+            return [
+                (base_lr - self.warmup_lr) * self.last_epoch / self.warmup_steps
+                + self.warmup_lr
+                for base_lr in self.base_lrs
+            ]
+        else:
+            current_steps = self.last_epoch - self.warmup_steps
+            return [
+                0.5
+                * base_lr
+                * (1 + math.cos(math.pi * current_steps / self.decay_steps))
+                for base_lr in self.base_lrs
+            ]

utils/metric.py

@@ -0,0 +1,58 @@
+from typing import List, Union
+import os
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+from torch.autograd import Variable
+import torch.optim as optim
+import numpy as np
+import torch
+__all__ = ["accuracy", "AverageMeter"]
+
+
+
+
+
+def accuracy(
+    output: torch.Tensor, target: torch.Tensor, topk=(1,)
+) -> List[torch.Tensor]:
+    """Computes the precision@k for the specified values of k."""
+    maxk = max(topk)
+    batch_size = target.shape[0]
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.reshape(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value.
+
+    Copied from: https://github.com/pytorch/examples/blob/master/imagenet/main.py
+    """
+
+    def __init__(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val: Union[torch.Tensor, np.ndarray, float, int], n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count

utils/misc.py

@@ -0,0 +1,122 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import yaml
+from torch.nn.modules.batchnorm import _BatchNorm
+
+__all__ = [
+    "make_divisible",
+    "load_state_dict_from_file",
+    "list_mean",
+    "list_sum",
+    "parse_unknown_args",
+    "partial_update_config",
+    "remove_bn",
+    "get_same_padding",
+    "torch_random_choices",
+]
+
+
+def make_divisible(
+    v: Union[int, float], divisor: Optional[int], min_val=None
+) -> Union[int, float]:
+    """This function is taken from the original tf repo.
+
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_val:
+    :return:
+    """
+    if divisor is None:
+        return v
+
+    if min_val is None:
+        min_val = divisor
+    new_v = max(min_val, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+def load_state_dict_from_file(file: str) -> Dict[str, torch.Tensor]:
+    checkpoint = torch.load(file, map_location="cpu")
+    if "state_dict" in checkpoint:
+        checkpoint = checkpoint["state_dict"]
+    return checkpoint
+
+
+def list_sum(x: List) -> Any:
+    return x[0] if len(x) == 1 else x[0] + list_sum(x[1:])
+
+
+def list_mean(x: List) -> Any:
+    return list_sum(x) / len(x)
+
+
+def parse_unknown_args(unknown: List) -> Dict:
+    """Parse unknown args."""
+    index = 0
+    parsed_dict = {}
+    while index < len(unknown):
+        key, val = unknown[index], unknown[index + 1]
+        index += 2
+        if key.startswith("--"):
+            key = key[2:]
+            try:
+                # try parsing with yaml
+                if "{" in val and "}" in val and ":" in val:
+                    val = val.replace(":", ": ")  # add space manually for dict
+                out_val = yaml.safe_load(val)
+            except ValueError:
+                # return raw string if parsing fails
+                out_val = val
+            parsed_dict[key] = out_val
+    return parsed_dict
+
+
+def partial_update_config(config: Dict, partial_config: Dict):
+    for key in partial_config:
+        if (
+            key in config
+            and isinstance(partial_config[key], Dict)
+            and isinstance(config[key], Dict)
+        ):
+            partial_update_config(config[key], partial_config[key])
+        else:
+            config[key] = partial_config[key]
+
+
+def remove_bn(model: nn.Module) -> None:
+    for m in model.modules():
+        if isinstance(m, _BatchNorm):
+            m.weight = m.bias = None
+            m.forward = lambda x: x
+
+
+def get_same_padding(kernel_size: Union[int, Tuple[int, int]]) -> Union[int, tuple]:
+    if isinstance(kernel_size, tuple):
+        assert len(kernel_size) == 2, f"invalid kernel size: {kernel_size}"
+        p1 = get_same_padding(kernel_size[0])
+        p2 = get_same_padding(kernel_size[1])
+        return p1, p2
+    else:
+        assert isinstance(
+            kernel_size, int
+        ), "kernel size should be either `int` or `tuple`"
+        assert kernel_size % 2 > 0, "kernel size should be odd number"
+        return kernel_size // 2
+
+
+def torch_random_choices(
+    src_list: List[Any],
+    generator: Optional[torch.Generator],
+    k=1,
+) -> Union[Any, List[Any]]:
+    rand_idx = torch.randint(low=0, high=len(src_list), generator=generator, size=(k,))
+    out_list = [src_list[i] for i in rand_idx]
+    return out_list[0] if k == 1 else out_list

utils/profile.py

@@ -0,0 +1,40 @@
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torchprofile import profile_macs
+
+__all__ = ["is_parallel", "get_module_device", "trainable_param_num", "inference_macs"]
+
+
+def is_parallel(model: nn.Module) -> bool:
+    return isinstance(
+        model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
+    )
+
+
+def get_module_device(module: nn.Module) -> torch.device:
+    return module.parameters().__next__().device
+
+
+def trainable_param_num(network: nn.Module, unit=1e6) -> float:
+    return sum(p.numel() for p in network.parameters() if p.requires_grad) / unit
+
+
+def inference_macs(
+    network: nn.Module,
+    args: Tuple = (),
+    data_shape: Optional[Tuple] = None,
+    unit: float = 1e6,
+) -> float:
+    if is_parallel(network):
+        network = network.module
+    if data_shape is not None:
+        if len(args) > 0:
+            raise ValueError("Please provide either data_shape or args tuple.")
+        args = (torch.zeros(data_shape, device=get_module_device(network)),)
+    is_training = network.training
+    network.eval()
+    macs = profile_macs(network, args=args) / unit
+    network.train(is_training)
+    return macs