| 88 | |
| 89 | |
| 90 | class TrainingAgent: |
| 91 | config: Dict[str, Any] = None |
| 92 | train_loader: DataLoader = None |
| 93 | test_loader: DataLoader = None |
| 94 | num_classes: int = None |
| 95 | network: torch.nn.Module = None |
| 96 | optimizer: torch.optim.Optimizer = None |
| 97 | scheduler: torch.optim.lr_scheduler.LRScheduler = None |
| 98 | criterion: torch.nn.Module = None |
| 99 | output_filename: Path = None |
| 100 | checkpoint = None |
| 101 | early_stop = None |
| 102 | gm = None |
| 103 | |
| 104 | def __init__( |
| 105 | self, |
| 106 | config_path: Path, |
| 107 | device: str, |
| 108 | output_path: Path, |
| 109 | data_path: Path, |
| 110 | checkpoint_path: Path, |
| 111 | resume: Path = None, |
| 112 | save_freq: int = 25, |
| 113 | dist: bool = False, |
| 114 | autocast: bool = False, |
| 115 | pretrained: bool = False) -> None: |
| 116 | |
| 117 | self.dist = dist |
| 118 | if self.dist: |
| 119 | self.gpu = int(os.environ["LOCAL_RANK"]) |
| 120 | else: |
| 121 | self.gpu = device |
| 122 | self.best_acc1 = 0. |
| 123 | self.start_epoch = 0 |
| 124 | self.start_trial = 0 |
| 125 | self.device = device |
| 126 | self.autocast = autocast |
| 127 | self.data_path = data_path |
| 128 | self.output_path = output_path |
| 129 | self.checkpoint_path = checkpoint_path |
| 130 | self.resume = resume |
| 131 | self.save_freq = save_freq |
| 132 | self.pretrained = pretrained |
| 133 | |
| 134 | self.load_config(config_path, data_path) |
| 135 | print("Experiment Configuration") |
| 136 | print("-" * 45) |
| 137 | for k, v in self.config.items(): |
| 138 | if isinstance(v, list) or isinstance(v, dict): |
| 139 | print(f" {k:<20} {v}") |
| 140 | else: |
| 141 | print(f" {k:<20} {v:<20}") |
| 142 | print("-" * 45) |
| 143 | |
| 144 | def load_config(self, config_path: Path, data_path: Path) -> None: |
| 145 | with config_path.open() as f: |
| 146 | self.config = config = parse_config( |
| 147 | yaml.load(f, Loader=yaml.Loader)) |