(cfg, local_rank, distributed)
| 35 | |
| 36 | |
| 37 | def train(cfg, local_rank, distributed): |
| 38 | logger = logging.getLogger("SourceOnly.trainer") |
| 39 | logger.info("Start training") |
| 40 | |
| 41 | feature_extractor = build_feature_extractor(cfg) |
| 42 | device = torch.device(cfg.MODEL.DEVICE) |
| 43 | feature_extractor.to(device) |
| 44 | |
| 45 | classifier = build_classifier(cfg) |
| 46 | classifier.to(device) |
| 47 | |
| 48 | if local_rank == 0: |
| 49 | print(feature_extractor) |
| 50 | print(classifier) |
| 51 | |
| 52 | batch_size = cfg.SOLVER.BATCH_SIZE # default: 8 |
| 53 | if distributed: |
| 54 | pg1 = torch.distributed.new_group(range(torch.distributed.get_world_size())) |
| 55 | batch_size = int(cfg.SOLVER.BATCH_SIZE / torch.distributed.get_world_size()) |
| 56 | if not cfg.MODEL.FREEZE_BN: |
| 57 | # if don't freeze BN, then replace BN with SyncBatchNorm |
| 58 | # for self-supervised learning, use SyncBatchNorm |
| 59 | feature_extractor = torch.nn.SyncBatchNorm.convert_sync_batchnorm(feature_extractor) |
| 60 | feature_extractor = torch.nn.parallel.DistributedDataParallel( |
| 61 | feature_extractor, device_ids=[local_rank], output_device=local_rank, |
| 62 | find_unused_parameters=True, process_group=pg1 |
| 63 | ) |
| 64 | pg2 = torch.distributed.new_group(range(torch.distributed.get_world_size())) |
| 65 | classifier = torch.nn.parallel.DistributedDataParallel( |
| 66 | classifier, device_ids=[local_rank], output_device=local_rank, |
| 67 | find_unused_parameters=True, process_group=pg2 |
| 68 | ) |
| 69 | torch.autograd.set_detect_anomaly(True) |
| 70 | torch.distributed.barrier() |
| 71 | |
| 72 | optimizer_fea = torch.optim.SGD(feature_extractor.parameters(), lr=cfg.SOLVER.BASE_LR, momentum=cfg.SOLVER.MOMENTUM, |
| 73 | weight_decay=cfg.SOLVER.WEIGHT_DECAY) |
| 74 | optimizer_fea.zero_grad() |
| 75 | |
| 76 | optimizer_cls = torch.optim.SGD(classifier.parameters(), lr=cfg.SOLVER.BASE_LR * 10, momentum=cfg.SOLVER.MOMENTUM, |
| 77 | weight_decay=cfg.SOLVER.WEIGHT_DECAY) |
| 78 | optimizer_cls.zero_grad() |
| 79 | |
| 80 | output_dir = cfg.OUTPUT_DIR |
| 81 | |
| 82 | save_to_disk = local_rank == 0 |
| 83 | |
| 84 | iteration = 0 |
| 85 | |
| 86 | if cfg.resume: |
| 87 | logger.info("Loading checkpoint from {}".format(cfg.resume)) |
| 88 | checkpoint = torch.load(cfg.resume, map_location=torch.device('cpu')) |
| 89 | model_weights = checkpoint['feature_extractor'] if distributed else strip_prefix_if_present( |
| 90 | checkpoint['feature_extractor'], 'module.') |
| 91 | feature_extractor.load_state_dict(model_weights) |
| 92 | classifier_weights = checkpoint['classifier'] if distributed else strip_prefix_if_present( |
| 93 | checkpoint['classifier'], 'module.') |
| 94 | classifier.load_state_dict(classifier_weights) |
no test coverage detected