(cfg, logger)
| 50 | |
| 51 | |
| 52 | def build_solver(cfg, logger): |
| 53 | task = core.Configurable.load_config_dict(cfg.task) |
| 54 | |
| 55 | # build solver |
| 56 | solver = DiffusionEngine(task, None, None, None, None, None, **cfg.engine) |
| 57 | if "checkpoint" in cfg: |
| 58 | solver.load(cfg.checkpoint, load_optimizer=cfg.get("load_optimizer", False)) |
| 59 | |
| 60 | if "model_checkpoint" in cfg: |
| 61 | model_checkpoint = os.path.expanduser(cfg.model_checkpoint) |
| 62 | model_dict = torch.load(model_checkpoint, map_location=torch.device('cpu'))["model"] |
| 63 | missing_keys, unexpected_keys = task.load_state_dict(model_dict, strict=False) |
| 64 | |
| 65 | # Calculate the parameter number of the model |
| 66 | if comm.get_rank() == 0: |
| 67 | logger.warning("#parameter: %d" % sum(p.numel() for p in task.parameters() if p.requires_grad)) |
| 68 | |
| 69 | return solver |
| 70 | |
| 71 | |
| 72 | if __name__ == "__main__": |
no test coverage detected