(dataloader, model, model_refine, optimizer, epoch)
| 29 | os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu |
| 30 | |
| 31 | def train(dataloader, model, model_refine, optimizer, epoch): |
| 32 | model.train() |
| 33 | loss_all = {'loss': AccumLoss()} |
| 34 | |
| 35 | for i, data in enumerate(tqdm(dataloader, 0)): |
| 36 | batch_cam, gt_3D, input_2D, input_2D_GT, action, subject, cam_ind = data |
| 37 | [input_2D, input_2D_GT, gt_3D, batch_cam] = get_varialbe('train', [input_2D, input_2D_GT, gt_3D, batch_cam]) |
| 38 | |
| 39 | output_3D = model(input_2D) |
| 40 | |
| 41 | out_target = gt_3D.clone() |
| 42 | out_target[:, :, args.root_joint] = 0 |
| 43 | out_target = out_target[:, args.pad].unsqueeze(1) |
| 44 | |
| 45 | if args.refine: |
| 46 | model_refine.train() |
| 47 | output_3D = refine_model(model_refine, output_3D, input_2D, gt_3D, batch_cam, args.pad, args.root_joint) |
| 48 | loss = eval_cal.mpjpe(output_3D, out_target) |
| 49 | else: |
| 50 | loss = eval_cal.mpjpe(output_3D, out_target) |
| 51 | |
| 52 | optimizer.zero_grad() |
| 53 | loss.backward() |
| 54 | optimizer.step() |
| 55 | |
| 56 | N = input_2D.shape[0] |
| 57 | loss_all['loss'].update(loss.detach().cpu().numpy() * N, N) |
| 58 | |
| 59 | return loss_all['loss'].avg |
| 60 | |
| 61 | |
| 62 | def test(actions, dataloader, model, model_refine): |
no test coverage detected