(config)
| 104 | |
| 105 | |
| 106 | def main(config): |
| 107 | _, dataset_val, data_loader_train, data_loader_val, mixup_fn = build_loader(config) |
| 108 | |
| 109 | logger.info(f"Creating model:{config.MODEL.TYPE}/{config.MODEL.NAME}") |
| 110 | model = build_model(config) |
| 111 | model.cuda() |
| 112 | logger.info(str(model)) |
| 113 | |
| 114 | # fix parameters except head |
| 115 | for name, p in model.named_parameters(): |
| 116 | if 'head' not in name: |
| 117 | p.requires_grad = False |
| 118 | |
| 119 | optimizer = build_optimizer(config, model) |
| 120 | if config.AMP_OPT_LEVEL != "O0": |
| 121 | model, optimizer = amp.initialize(model, optimizer, opt_level=config.AMP_OPT_LEVEL) |
| 122 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[config.LOCAL_RANK], broadcast_buffers=False) |
| 123 | model_without_ddp = model.module |
| 124 | |
| 125 | # load self-supervised pre-trained model |
| 126 | load_pretrained(model_without_ddp, config.LINEAR_EVAL.PRETRAINED, logger) |
| 127 | |
| 128 | n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| 129 | logger.info(f"number of params: {n_parameters}") |
| 130 | if hasattr(model_without_ddp, 'flops'): |
| 131 | flops = model_without_ddp.flops() |
| 132 | logger.info(f"number of GFLOPs: {flops / 1e9}") |
| 133 | |
| 134 | lr_scheduler = build_scheduler(config, optimizer, len(data_loader_train)) |
| 135 | |
| 136 | if config.AUG.MIXUP > 0.: |
| 137 | # smoothing is handled with mixup label transform |
| 138 | criterion = SoftTargetCrossEntropy() |
| 139 | elif config.MODEL.LABEL_SMOOTHING > 0.: |
| 140 | criterion = LabelSmoothingCrossEntropy(smoothing=config.MODEL.LABEL_SMOOTHING) |
| 141 | else: |
| 142 | criterion = torch.nn.CrossEntropyLoss() |
| 143 | |
| 144 | max_accuracy = 0.0 |
| 145 | |
| 146 | if config.TRAIN.AUTO_RESUME: |
| 147 | resume_file = auto_resume_helper(config.OUTPUT) |
| 148 | if resume_file: |
| 149 | if config.MODEL.RESUME: |
| 150 | logger.warning(f"auto-resume changing resume file from {config.MODEL.RESUME} to {resume_file}") |
| 151 | config.defrost() |
| 152 | config.MODEL.RESUME = resume_file |
| 153 | config.freeze() |
| 154 | logger.info(f'auto resuming from {resume_file}') |
| 155 | else: |
| 156 | logger.info(f'no checkpoint found in {config.OUTPUT}, ignoring auto resume') |
| 157 | |
| 158 | if config.MODEL.RESUME: |
| 159 | max_accuracy = load_checkpoint(config, model_without_ddp, optimizer, lr_scheduler, logger) |
| 160 | acc1, acc5, loss = validate(config, data_loader_val, model) |
| 161 | logger.info(f"Accuracy of the network on the {len(dataset_val)} test images: {acc1:.1f}%") |
| 162 | if config.EVAL_MODE: |
| 163 | return |
no test coverage detected