| 82 | |
| 83 | |
| 84 | def build_callback(name, kwargs): |
| 85 | if name == "lr_monitor": |
| 86 | return LRMonitor() |
| 87 | elif name == "memory_monitor": |
| 88 | return MemoryMonitor() |
| 89 | elif name == "speed_monitor": |
| 90 | return SpeedMonitor( |
| 91 | window_size=kwargs.get("window_size", 1), |
| 92 | gpu_flops_available=kwargs.get("gpu_flops_available", None), |
| 93 | ) |
| 94 | elif name == "runtime_estimator": |
| 95 | return RuntimeEstimator() |
| 96 | elif name == "optimizer_monitor": |
| 97 | return OptimizerMonitor( |
| 98 | log_optimizer_metrics=kwargs.get("log_optimizer_metrics", True), |
| 99 | ) |
| 100 | else: |
| 101 | raise ValueError(f"Not sure how to build callback: {name}") |
| 102 | |
| 103 | |
| 104 | def build_logger(name, kwargs): |