| 30 | |
| 31 | |
| 32 | class Trainer(DefaultTrainer): |
| 33 | # """ |
| 34 | # Extension of the Trainer class adapted to SparseRCNN. |
| 35 | # """ |
| 36 | |
| 37 | @classmethod |
| 38 | def build_evaluator(cls, cfg, dataset_name, output_folder=None): |
| 39 | """ |
| 40 | Create evaluator(s) for a given dataset. |
| 41 | This uses the special metadata "evaluator_type" associated with each builtin dataset. |
| 42 | For your own dataset, you can simply create an evaluator manually in your |
| 43 | script and do not have to worry about the hacky if-else logic here. |
| 44 | """ |
| 45 | if output_folder is None: |
| 46 | output_folder = os.path.join(cfg.OUTPUT_DIR, "inference") |
| 47 | return COCOEvaluator(dataset_name, cfg, True, output_folder) |
| 48 | |
| 49 | @classmethod |
| 50 | def build_train_loader(cls, cfg): |
| 51 | mapper = SparseRCNNDatasetMapper(cfg, is_train=True) |
| 52 | return build_detection_train_loader(cfg, mapper=mapper) |
| 53 | |
| 54 | @classmethod |
| 55 | def build_optimizer(cls, cfg, model): |
| 56 | params: List[Dict[str, Any]] = [] |
| 57 | memo: Set[torch.nn.parameter.Parameter] = set() |
| 58 | for key, value in model.named_parameters(recurse=True): |
| 59 | if not value.requires_grad: |
| 60 | continue |
| 61 | # Avoid duplicating parameters |
| 62 | if value in memo: |
| 63 | continue |
| 64 | memo.add(value) |
| 65 | lr = cfg.SOLVER.BASE_LR |
| 66 | weight_decay = cfg.SOLVER.WEIGHT_DECAY |
| 67 | if "backbone" in key: |
| 68 | lr = lr * cfg.SOLVER.BACKBONE_MULTIPLIER |
| 69 | params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}] |
| 70 | |
| 71 | def maybe_add_full_model_gradient_clipping(optim): # optim: the optimizer class |
| 72 | # detectron2 doesn't have full model gradient clipping now |
| 73 | clip_norm_val = cfg.SOLVER.CLIP_GRADIENTS.CLIP_VALUE |
| 74 | enable = ( |
| 75 | cfg.SOLVER.CLIP_GRADIENTS.ENABLED |
| 76 | and cfg.SOLVER.CLIP_GRADIENTS.CLIP_TYPE == "full_model" |
| 77 | and clip_norm_val > 0.0 |
| 78 | ) |
| 79 | |
| 80 | class FullModelGradientClippingOptimizer(optim): |
| 81 | def step(self, closure=None): |
| 82 | all_params = itertools.chain(*[x["params"] for x in self.param_groups]) |
| 83 | torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val) |
| 84 | super().step(closure=closure) |
| 85 | |
| 86 | return FullModelGradientClippingOptimizer if enable else optim |
| 87 | |
| 88 | optimizer_type = cfg.SOLVER.OPTIMIZER |
| 89 | if optimizer_type == "SGD": |