(cfg)
| 54 | return sum(p.numel() for p in model.parameters() if p.requires_grad) |
| 55 | |
| 56 | def train(cfg): |
| 57 | model = nn.DataParallel(build_flowformer(cfg)) |
| 58 | loguru_logger.info("Parameter Count: %d" % count_parameters(model)) |
| 59 | |
| 60 | if cfg.restore_ckpt is not None: |
| 61 | print("[Loading ckpt from {}]".format(cfg.restore_ckpt)) |
| 62 | model.load_state_dict(torch.load(cfg.restore_ckpt), strict=True) |
| 63 | |
| 64 | model.cuda() |
| 65 | model.train() |
| 66 | |
| 67 | train_loader = datasets.fetch_dataloader(cfg) |
| 68 | optimizer, scheduler = fetch_optimizer(model, cfg.trainer) |
| 69 | |
| 70 | total_steps = 0 |
| 71 | scaler = GradScaler(enabled=cfg.mixed_precision) |
| 72 | logger = Logger(model, scheduler, cfg) |
| 73 | |
| 74 | add_noise = False |
| 75 | |
| 76 | should_keep_training = True |
| 77 | while should_keep_training: |
| 78 | |
| 79 | for i_batch, data_blob in enumerate(train_loader): |
| 80 | optimizer.zero_grad() |
| 81 | image1, image2, flow, valid = [x.cuda() for x in data_blob] |
| 82 | |
| 83 | if cfg.add_noise: |
| 84 | stdv = np.random.uniform(0.0, 5.0) |
| 85 | image1 = (image1 + stdv * torch.randn(*image1.shape).cuda()).clamp(0.0, 255.0) |
| 86 | image2 = (image2 + stdv * torch.randn(*image2.shape).cuda()).clamp(0.0, 255.0) |
| 87 | |
| 88 | output = {} |
| 89 | flow_predictions = model(image1, image2, output) |
| 90 | loss, metrics = sequence_loss(flow_predictions, flow, valid, cfg) |
| 91 | scaler.scale(loss).backward() |
| 92 | scaler.unscale_(optimizer) |
| 93 | torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.trainer.clip) |
| 94 | |
| 95 | scaler.step(optimizer) |
| 96 | scheduler.step() |
| 97 | scaler.update() |
| 98 | |
| 99 | metrics.update(output) |
| 100 | logger.push(metrics) |
| 101 | |
| 102 | ### change evaluate to functions |
| 103 | |
| 104 | if total_steps % cfg.val_freq == cfg.val_freq - 1: |
| 105 | PATH = '%s/%d_%s.pth' % (cfg.log_dir, total_steps+1, cfg.name) |
| 106 | # torch.save(model.state_dict(), PATH) |
| 107 | |
| 108 | results = {} |
| 109 | for val_dataset in cfg.validation: |
| 110 | if val_dataset == 'chairs': |
| 111 | results.update(evaluate.validate_chairs(model.module)) |
| 112 | elif val_dataset == 'sintel': |
| 113 | results.update(evaluate.validate_sintel(model.module)) |
no test coverage detected