def _samplers_loader(r: Registry):
from torch.utils.data import sampler as s
factories = {
k: v
for k, v in s.__dict__.items() if "Sampler" in k and k != "Sampler"
}
r.add(**factories)
from catalyst.data import sampler
r.add_from_module(sampler)
def _optimizers_loader(r: Registry):
from catalyst.contrib import optimizers as m
r.add_from_module(m)
def _schedulers_loader(r: Registry):
from catalyst.contrib import schedulers as m
r.add_from_module(m)
def _modules_loader(r: Registry):
from catalyst.contrib import modules as m
r.add_from_module(m)
def _criterion_loader(r: Registry):
from catalyst.contrib import criterion as m
r.add_from_module(m)
def _grad_clip_loader(r: Registry):
from torch.nn.utils import clip_grad as m
r.add_from_module(m)
class _GradClipperWrap:
def __init__(self, fn, args, kwargs):
self.fn = fn
self.args = args
self.kwargs = kwargs
def __call__(self, x):
self.fn(x, *self.args, **self.kwargs)
def _grad_clip_loader(r: Registry):
from torch.nn.utils import clip_grad as m
r.add_from_module(m)
GRAD_CLIPPERS = Registry("func", default_meta_factory=_GradClipperWrap)
GRAD_CLIPPERS.late_add(_grad_clip_loader)
def _criterion_loader(r: Registry):
from catalyst.contrib import criterion as m
r.add_from_module(m)
CRITERIONS = Registry("criterion")
CRITERIONS.late_add(_criterion_loader)
Criterion = CRITERIONS.add
def _model_loader(r: Registry):
from catalyst.contrib import models as m
def _schedulers_loader(r: Registry):
from torch.optim import lr_scheduler as m
r.add_from_module(m)