| 16 | #! Traning # |
| 17 | #!=======================================================================# |
| 18 | class LcdNet(object): |
| 19 | def __init__(self, config, logger, use_cuda, device, neptune=None, gpu_ids=None): |
| 20 | super(LcdNet, self).__init__() |
| 21 | |
| 22 | self.neptune = neptune |
| 23 | self.config = config |
| 24 | self.logger = logger |
| 25 | |
| 26 | #! Let's define different models |
| 27 | self.model = make_models(config) |
| 28 | log_print("Runing {}".format(config.MODEL.NAME), 'r') |
| 29 | log_print("#params of model: {}".format(sum([x.numel() |
| 30 | for x in self.model.parameters()])), 'b') |
| 31 | log_print("#trainable params of model: {}".format(sum([x.numel() |
| 32 | for x in self.model.parameters() if x.requires_grad])), 'b') |
| 33 | |
| 34 | #! Data Parallel |
| 35 | if use_cuda: |
| 36 | if gpu_ids is not None and len(gpu_ids)>1 : |
| 37 | print("Let's use", len(gpu_ids), "GPUs!") |
| 38 | self.model = nn.DataParallel(self.model, device_ids=gpu_ids) |
| 39 | self.model = self.model.to(device) |
| 40 | |
| 41 | #! Loss Function |
| 42 | self.criterion = make_losses(config).to(device) |
| 43 | |
| 44 | #! Optimizer and Scheduler |
| 45 | if config.TRAINING.OPTIMIZER.NAME == 'Adam': |
| 46 | self.optimizer = torch.optim.Adam( |
| 47 | filter(lambda p: p.requires_grad, self.model.parameters()), |
| 48 | lr=config.TRAINING.OPTIMIZER.INIT_LEARNING_RATE) |
| 49 | else: |
| 50 | raise NotImplementedError(f"Unrecognized Optimizer {config.TRAINING.OPTIMIZER.NAME}") |
| 51 | |
| 52 | if config.TRAINING.SCHEDULER.NAME == None: |
| 53 | self.scheduler = None |
| 54 | elif config.TRAINING.SCHEDULER.NAME == "StepLR": |
| 55 | self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, |
| 56 | step_size=config.TRAINING.SCHEDULER.STEP_SIZE, |
| 57 | gamma=config.TRAINING.SCHEDULER.GAMMA) |
| 58 | elif config.TRAINING.SCHEDULER.NAME == "MultiStepLR": |
| 59 | self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, |
| 60 | milestones=config.TRAINING.SCHEDULER.MILESTONES, |
| 61 | gamma=config.TRAINING.SCHEDULER.GAMMA) |
| 62 | else: |
| 63 | raise NotImplementedError(f"Unrecognized Scheduler {config.TRAINING.SCHEDULER.NAME}") |
| 64 | |
| 65 | #! Resume Training & Testing |
| 66 | if config.TRAINING.IS_TRAIN: |
| 67 | if config.TRAINING.RESUME: |
| 68 | self.epoch = self.load_checkpoint(config.WEIGHT.LOAD_ADDRESS, config.TRAINING.RESUME) + 1 |
| 69 | else: |
| 70 | self.epoch = 1 |
| 71 | else: |
| 72 | self.load_checkpoint(config.WEIGHT.LOAD_ADDRESS) |
| 73 | |
| 74 | def train_lcd(self, x): |
| 75 | """[summary] |