(self, cfg, mode='train', task='rec')
| 36 | class Trainer(object): |
| 37 | |
| 38 | def __init__(self, cfg, mode='train', task='rec'): |
| 39 | self.cfg = cfg.cfg |
| 40 | self.task = task |
| 41 | self.local_rank = (int(os.environ['LOCAL_RANK']) |
| 42 | if 'LOCAL_RANK' in os.environ else 0) |
| 43 | self.set_device(self.cfg['Global']['device']) |
| 44 | mode = mode.lower() |
| 45 | assert mode in [ |
| 46 | 'train_eval', |
| 47 | 'train', |
| 48 | 'eval', |
| 49 | 'test', |
| 50 | ], 'mode should be train, eval and test' |
| 51 | if torch.cuda.device_count() > 1 and 'train' in mode: |
| 52 | torch.distributed.init_process_group(backend='nccl') |
| 53 | torch.cuda.set_device(self.device) |
| 54 | self.cfg['Global']['distributed'] = True |
| 55 | else: |
| 56 | self.cfg['Global']['distributed'] = False |
| 57 | self.local_rank = 0 |
| 58 | |
| 59 | self.cfg['Global']['output_dir'] = self.cfg['Global'].get( |
| 60 | 'output_dir', 'output') |
| 61 | os.makedirs(self.cfg['Global']['output_dir'], exist_ok=True) |
| 62 | |
| 63 | self.writer = None |
| 64 | if is_main_process( |
| 65 | ) and self.cfg['Global']['use_tensorboard'] and 'train' in mode: |
| 66 | import wandb |
| 67 | from torch.utils.tensorboard import SummaryWriter |
| 68 | wandb.init(project='demo-sync-tb', |
| 69 | name=self.cfg['Global'].get('run_name', |
| 70 | 'log_wandb_openocr'), |
| 71 | sync_tensorboard=True) |
| 72 | |
| 73 | self.writer = SummaryWriter(self.cfg['Global']['output_dir']) |
| 74 | |
| 75 | self.logger = get_logger( |
| 76 | 'openrec' if task == 'rec' else 'opendet', |
| 77 | os.path.join(self.cfg['Global']['output_dir'], 'train.log') |
| 78 | if 'train' in mode else None, |
| 79 | ) |
| 80 | |
| 81 | cfg.print_cfg(self.logger.info) |
| 82 | |
| 83 | if self.cfg['Global']['device'] == 'gpu' and self.device.type == 'cpu': |
| 84 | self.logger.info('cuda is not available, auto switch to cpu') |
| 85 | |
| 86 | self.set_random_seed(self.cfg['Global'].get('seed', 48)) |
| 87 | |
| 88 | # build data loader |
| 89 | self.train_dataloader = None |
| 90 | if 'train' in mode: |
| 91 | if is_main_process(): |
| 92 | cfg.save( |
| 93 | os.path.join(self.cfg['Global']['output_dir'], |
| 94 | 'config.yml'), self.cfg) |
| 95 | self.train_dataloader = build_dataloader(self.cfg, |
nothing calls this directly
no test coverage detected