def build_recognizer(cfg, device): world_size = du.get_world_size() model = registry.RECOGNIZER[cfg.MODEL.RECOGNIZER.TYPE](cfg) if cfg.MODEL.NORM.SYNC_BN and world_size > 1: logger.info("start sync BN on the process group of {}".format( du.LOCAL_RANK_GROUP)) convert_sync_bn(model, du.LOCAL_PROCESS_GROUP) preloaded = cfg.MODEL.RECOGNIZER.PRELOADED if preloaded != "": logger.info(f'load pretrained: {preloaded}') check_pointer = CheckPointer(model) check_pointer.load(preloaded, map_location=device) logger.info("finish loading model weights") if cfg.MODEL.CONV.ADD_BLOCKS is not None: assert isinstance(cfg.MODEL.CONV.ADD_BLOCKS, tuple) for add_block in cfg.MODEL.CONV.ADD_BLOCKS: if add_block == 'RepVGGBlock': insert_repvgg_block(model) if add_block == 'ACBlock': insert_acblock(model) model = model.to(device=device) if du.get_world_size() > 1: model = DDP(model, device_ids=[device], output_device=device, find_unused_parameters=True) return model
def build_dataloader(cfg, dataset, is_train=True): world_size = du.get_world_size() num_gpus = cfg.NUM_GPUS rank = du.get_rank() if is_train: batch_size = cfg.DATALOADER.TRAIN_BATCH_SIZE if num_gpus > 1: sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) else: sampler = RandomSampler(dataset) else: batch_size = cfg.DATALOADER.TEST_BATCH_SIZE if num_gpus > 1: sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False) else: sampler = SequentialSampler(dataset) data_loader = DataLoader( dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, sampler=sampler, batch_size=batch_size, drop_last=False, # [When to set pin_memory to true?](https://discuss.pytorch.org/t/when-to-set-pin-memory-to-true/19723) pin_memory=True) return data_loader
def setup_logging(output_dir=None): """ Sets up the logging for multiple processes. Only enable the logging for the master process, and suppress logging for the non-master processes. """ # Set up logging format. _FORMAT = "[%(levelname)s: %(filename)s: %(lineno)4d]: %(message)s" if du.is_master_proc(): # Enable logging for the master process. logging.root.handlers = [] else: # Suppress logging for non-master processes. _suppress_print() logger = logging.getLogger() logger.setLevel(logging.DEBUG) logger.propagate = False plain_formatter = logging.Formatter( "[%(asctime)s][%(levelname)s] %(filename)s: %(lineno)3d: %(message)s", datefmt="%m/%d %H:%M:%S", ) if du.is_master_proc(): ch = logging.StreamHandler(stream=sys.stdout) ch.setLevel(logging.DEBUG) ch.setFormatter(plain_formatter) logger.addHandler(ch) if output_dir is not None and du.is_master_proc(du.get_world_size()): filename = os.path.join(output_dir, "stdout.log") fh = logging.FileHandler(filename) fh.setLevel(logging.DEBUG) fh.setFormatter(plain_formatter) logger.addHandler(fh)
def build_sampler(cfg, dataset): world_size = du.get_world_size() num_gpus = cfg.NUM_GPUS rank = du.get_rank() if num_gpus > 1: shuffle = cfg.DATALOADER.RANDOM_SAMPLE sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=shuffle) elif cfg.DATALOADER.RANDOM_SAMPLE: sampler = RandomSampler(dataset) else: sampler = SequentialSampler(dataset) return sampler
def build_dataloader(cfg, is_train=True): transform = build_transform(cfg, is_train=is_train) dataset = build_dataset(cfg, transform=transform, is_train=is_train) world_size = du.get_world_size() num_gpus = cfg.NUM_GPUS rank = du.get_rank() if is_train: batch_size = cfg.DATALOADER.TRAIN_BATCH_SIZE drop_last = True if num_gpus > 1: sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) else: sampler = RandomSampler(dataset) else: batch_size = cfg.DATALOADER.TEST_BATCH_SIZE drop_last = False if num_gpus > 1: sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False) else: sampler = SequentialSampler(dataset) data_loader = DataLoader(dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, sampler=sampler, batch_size=batch_size, drop_last=drop_last, pin_memory=True) return data_loader