| 177 | cleanup() |
| 178 | |
| 179 | def config_prase(args): |
| 180 | torch.manual_seed(args.seed) |
| 181 | np.random.seed(args.seed) |
| 182 | random.seed(args.seed) |
| 183 | |
| 184 | if args.resume is None: |
| 185 | with open(args.config, "r") as f: |
| 186 | config = yaml.safe_load(f) |
| 187 | |
| 188 | tmp_config_path = os.path.join(os.path.dirname(args.config), "params_tmp.yaml") |
| 189 | |
| 190 | config.update({'resume' : args.resume}) |
| 191 | config.update({'seed' : args.seed}) |
| 192 | config.update({'debug' : args.debug}) |
| 193 | config.update({'world_size' : torch.cuda.device_count()}) |
| 194 | |
| 195 | if not os.path.exists(config["workspace"]): |
| 196 | os.mkdir(config["workspace"]) |
| 197 | |
| 198 | config['timestamp'] = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime()) |
| 199 | workspace = os.path.join(config["workspace"], config['timestamp']) |
| 200 | if not os.path.exists(workspace): |
| 201 | os.mkdir(workspace) |
| 202 | config["local_workspace"] = workspace |
| 203 | |
| 204 | with open(tmp_config_path, "w") as f: |
| 205 | print("Dumping extra config file...") |
| 206 | yaml.dump(config, f) |
| 207 | |
| 208 | shutil.copy(tmp_config_path, os.path.join(workspace, "params.yaml")) |
| 209 | else: |
| 210 | config_path = os.path.join(args.workspace, args.resume, "params.yaml") |
| 211 | with open(config_path, "r") as f: |
| 212 | config = yaml.safe_load(f) |
| 213 | config.update({'resume' : args.resume}) |
| 214 | |
| 215 | config.update({'experiment_name': ''}) |
| 216 | return EasyDict(config) |
| 217 | |
| 218 | if __name__ == '__main__': |
| 219 | parser = argparse.ArgumentParser() |