73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import importlib
|
|
import argparse
|
|
import torch
|
|
import torch.nn as nn
|
|
import collections.abc
|
|
from itertools import repeat
|
|
|
|
def get_obj_from_str(string, reload=False):
|
|
module, cls = string.rsplit(".", 1)
|
|
if reload:
|
|
module_imp = importlib.import_module(module)
|
|
importlib.reload(module_imp)
|
|
return getattr(importlib.import_module(module, package=None), cls)
|
|
|
|
def instantiate_from_config(config):
|
|
if not isinstance(config, dict) or not "target" in config:
|
|
if config == '__is_first_stage__':
|
|
return None
|
|
elif config == "__is_unconditional__":
|
|
return None
|
|
raise KeyError("Expected key `target` to instantiate.")
|
|
return get_obj_from_str(config["target"])(**config.get("params", dict()))
|
|
|
|
def dict2namespace(config):
|
|
namespace = argparse.Namespace()
|
|
for key, value in config.items():
|
|
if isinstance(value, dict):
|
|
new_value = dict2namespace(value)
|
|
else:
|
|
new_value = value
|
|
setattr(namespace, key, new_value)
|
|
return namespace
|
|
|
|
# --- Native Replacements for timm utilities ---
|
|
|
|
def to_2tuple(x):
|
|
if isinstance(x, collections.abc.Iterable):
|
|
return x
|
|
return tuple(repeat(x, 2))
|
|
|
|
def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
|
|
# Simple implementation of truncated normal
|
|
with torch.no_grad():
|
|
size = tensor.shape
|
|
tmp = tensor.new_empty(size + (4,)).normal_()
|
|
valid = (tmp < b) & (tmp > a)
|
|
ind = valid.max(-1, keepdim=True)[1]
|
|
tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
|
|
tensor.data.mul_(std).add_(mean)
|
|
return tensor
|
|
|
|
def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
|
|
if drop_prob == 0. or not training:
|
|
return x
|
|
keep_prob = 1 - drop_prob
|
|
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
|
|
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
|
if keep_prob > 0.0 and scale_by_keep:
|
|
random_tensor.div_(keep_prob)
|
|
return x * random_tensor
|
|
|
|
class DropPath(nn.Module):
|
|
def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
|
|
super(DropPath, self).__init__()
|
|
self.drop_prob = drop_prob
|
|
self.scale_by_keep = scale_by_keep
|
|
|
|
def forward(self, x):
|
|
return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
|
|
|
|
def extra_repr(self):
|
|
return f'drop_prob={round(self.drop_prob,3):0.3f}'
|