Main api for training model.
(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
device='cuda',
meta=None)
| 33 | |
| 34 | |
| 35 | def train_model(model, |
| 36 | dataset, |
| 37 | cfg, |
| 38 | distributed=False, |
| 39 | validate=False, |
| 40 | timestamp=None, |
| 41 | device='cuda', |
| 42 | meta=None): |
| 43 | """Main api for training model.""" |
| 44 | logger = get_root_logger(cfg.log_level) |
| 45 | |
| 46 | # prepare data loaders |
| 47 | dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset] |
| 48 | |
| 49 | data_loaders = [ |
| 50 | build_dataloader( |
| 51 | ds, |
| 52 | cfg.data.samples_per_gpu, |
| 53 | cfg.data.workers_per_gpu, |
| 54 | # cfg.gpus will be ignored if distributed |
| 55 | num_gpus=len(cfg.gpu_ids), |
| 56 | dist=distributed, |
| 57 | round_up=True, |
| 58 | seed=cfg.seed) for ds in dataset |
| 59 | ] |
| 60 | |
| 61 | # determine whether use adversarial training precess or not |
| 62 | use_adverserial_train = cfg.get('use_adversarial_train', False) |
| 63 | |
| 64 | # put model on gpus |
| 65 | if distributed: |
| 66 | find_unused_parameters = cfg.get('find_unused_parameters', True) |
| 67 | # Sets the `find_unused_parameters` parameter in |
| 68 | # torch.nn.parallel.DistributedDataParallel |
| 69 | if use_adverserial_train: |
| 70 | # Use DistributedDataParallelWrapper for adversarial training |
| 71 | model = DistributedDataParallelWrapper( |
| 72 | model, |
| 73 | device_ids=[torch.cuda.current_device()], |
| 74 | broadcast_buffers=False, |
| 75 | find_unused_parameters=find_unused_parameters) |
| 76 | else: |
| 77 | model = MMDistributedDataParallel( |
| 78 | model.cuda(), |
| 79 | device_ids=[torch.cuda.current_device()], |
| 80 | broadcast_buffers=False, |
| 81 | find_unused_parameters=find_unused_parameters) |
| 82 | else: |
| 83 | if device == 'cuda': |
| 84 | model = MMDataParallel(model.cuda(cfg.gpu_ids[0]), |
| 85 | device_ids=cfg.gpu_ids) |
| 86 | elif device == 'cpu': |
| 87 | model = model.cpu() |
| 88 | else: |
| 89 | raise ValueError(F'unsupported device name {device}.') |
| 90 | |
| 91 | # build runner |
| 92 | optimizer = build_optimizers(model, cfg.optimizer) |
no test coverage detected