Update xmisc.scheduler/sampler

xautodl/xmisc/__init__.py

@@ -10,3 +10,23 @@ from .yaml_utils import load_yaml
 from .torch_utils import count_parameters
 
 from .logger_utils import Logger
+
+# sampler
+from .sampler_utils import BatchSampler
+
+# scheduler related
+from .scheduler_utils import CosineParamScheduler, WarmupParamScheduler, LRMultiplier
+
+
+def get_scheduler(indicator, lr):
+    if indicator == "warm-cos":
+        multiplier = WarmupParamScheduler(
+            CosineParamScheduler(lr, lr * 1e-3),
+            warmup_factor=0.001,
+            warmup_length=0.05,
+            warmup_method="linear",
+        )
+
+    else:
+        raise ValueError("Unknown indicator: {:}".format(indicator))
+    return multiplier