(model_wo_ddp)
| 196 | |
| 197 | |
| 198 | def log_model_statistic(model_wo_ddp): |
| 199 | n_parameters = sum(p.numel() for p in model_wo_ddp.parameters() |
| 200 | if p.requires_grad) |
| 201 | logger.info(f'number of params: {n_parameters}') |
| 202 | if hasattr(model_wo_ddp, 'flops'): |
| 203 | flops = model_wo_ddp.flops() |
| 204 | logger.info(f'number of GFLOPs: {flops / 1e9}') |
| 205 | |
| 206 | |
| 207 | def train_epoch(*, model, optimizer, data_loader, scheduler, criterion, mixup_fn, |