def test_add_module():
r = Registry("")
r.add_from_module(module)
r.get("foo")
with pytest.raises(RegistryException):
r.get_instance("bar")
def _model_loader(r: Registry):
from catalyst_rl.contrib import models as m
r.add_from_module(m)
try:
import segmentation_models_pytorch as smp
r.add_from_module(smp, prefix="smp.")
except ImportError:
pass
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_rl.data import sampler
r.add_from_module(sampler)
def _transforms_loader(r: Registry):
try:
import albumentations as m
r.add_from_module(m, prefix=["A.", "albu.", "albumentations."])
from albumentations import pytorch as p
r.add_from_module(p, prefix=["A.", "albu.", "albumentations."])
from catalyst_rl.contrib.data.cv import transforms as t
r.add_from_module(t, prefix=["catalyst_rl.", "C."])
except ImportError as ex:
if os.environ.get("USE_ALBUMENTATIONS", "0") == "1":
logger.warning(
"albumentations not available, to install albumentations, "
"run `pip install albumentations`.")
raise ex
def _dbs_late_add(r: Registry):
from . import db as m
r.add_from_module(m)
def _exploration_late_add(r: Registry):
from . import exploration as m
r.add_from_module(m)
def _env_late_add(r: Registry):
from . import environment as m
r.add_from_module(m)
def _onpolicy_algorithms_late_add(r: Registry):
from .onpolicy import algorithms as m
r.add_from_module(m)
def _agents_late_add(r: Registry):
from . import agent as m
r.add_from_module(m)
def _callbacks_loader(r: Registry):
from catalyst_rl.core import callbacks as m
r.add_from_module(m)
def _criterion_loader(r: Registry):
from catalyst_rl.contrib.nn import criterion as m
r.add_from_module(m)
def _modules_loader(r: Registry):
from catalyst_rl.contrib.nn import modules 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)
def _optimizers_loader(r: Registry):
from catalyst_rl.contrib.nn import optimizers as m
r.add_from_module(m)