Train the model.
(model, optimizer, lr_scheduler,
train_data, val_data, timers, args,
summary_writer=None, hooks={})
| 329 | |
| 330 | |
| 331 | def train(model, optimizer, lr_scheduler, |
| 332 | train_data, val_data, timers, args, |
| 333 | summary_writer=None, hooks={}): |
| 334 | """Train the model.""" |
| 335 | if train_data is not None: |
| 336 | train_data_iterator = iter(train_data) |
| 337 | else: |
| 338 | train_data_iterator = None |
| 339 | if val_data is not None: |
| 340 | val_data_iterator = iter(val_data) |
| 341 | else: |
| 342 | val_data_iterator = None |
| 343 | |
| 344 | # Turn on training mode which enables dropout. |
| 345 | model.train() |
| 346 | |
| 347 | # Tracking loss. |
| 348 | total_lm_loss = 0.0 |
| 349 | total_metrics = defaultdict(float) |
| 350 | total_metrics_cnt = defaultdict(int) |
| 351 | |
| 352 | # Iterations. |
| 353 | skipped_iters = 0 |
| 354 | |
| 355 | timers('interval time').start() |
| 356 | report_memory_flag = True |
| 357 | while args.iteration < args.train_iters: |
| 358 | if args.profiling != -1 and args.iteration == args.profiling: |
| 359 | torch.cuda.cudart().cudaProfilerStart() |
| 360 | |
| 361 | if args.profiling != -1 and args.iteration >= args.profiling: |
| 362 | torch.cuda.nvtx.range_push("iteration{}".format(args.iteration)) |
| 363 | lm_loss, skipped_iter, metrics = train_step(train_data_iterator, |
| 364 | model, |
| 365 | optimizer, |
| 366 | lr_scheduler, |
| 367 | args, timers, hooks=hooks) |
| 368 | skipped_iters += skipped_iter |
| 369 | if args.profiling != -1 and args.iteration >= args.profiling: |
| 370 | torch.cuda.nvtx.range_pop() |
| 371 | args.iteration += 1 |
| 372 | # Update losses. |
| 373 | total_lm_loss += lm_loss.data.detach().float() |
| 374 | for name in metrics: |
| 375 | if not 'eval' in name: |
| 376 | assert len(metrics[name].shape)==0, 'metrics without eval must be scalar' |
| 377 | value = metrics[name].data.detach().float().item() |
| 378 | if value > -99: |
| 379 | total_metrics[name] += value |
| 380 | total_metrics_cnt[name] += 1 |
| 381 | |
| 382 | # Logging. |
| 383 | if args.iteration % args.log_interval == 0: |
| 384 | learning_rate = optimizer.param_groups[0]['lr'] |
| 385 | avg_lm_loss = total_lm_loss.item() / args.log_interval |
| 386 | # average img & txt loss |
| 387 | avg_metrics = {} |
| 388 | for key in total_metrics: |
no test coverage detected