(opt, modelG, modelD, data_loader)
| 102 | return modelG, modelD, flowNet, optimizer_G, optimizer_D, optimizer_D_T |
| 103 | |
| 104 | def init_params(opt, modelG, modelD, data_loader): |
| 105 | iter_path = os.path.join(opt.checkpoints_dir, opt.name, 'iter.txt') |
| 106 | start_epoch, epoch_iter = 1, 0 |
| 107 | ### if continue training, recover previous states |
| 108 | if opt.continue_train: |
| 109 | if os.path.exists(iter_path): |
| 110 | start_epoch, epoch_iter = np.loadtxt(iter_path , delimiter=',', dtype=int) |
| 111 | print('Resuming from epoch %d at iteration %d' % (start_epoch, epoch_iter)) |
| 112 | if start_epoch > opt.niter: |
| 113 | modelG.module.update_learning_rate(start_epoch-1, 'G') |
| 114 | modelD.module.update_learning_rate(start_epoch-1, 'D') |
| 115 | if (opt.n_scales_spatial > 1) and (opt.niter_fix_global != 0) and (start_epoch > opt.niter_fix_global): |
| 116 | modelG.module.update_fixed_params() |
| 117 | if start_epoch > opt.niter_step: |
| 118 | data_loader.dataset.update_training_batch((start_epoch-1)//opt.niter_step) |
| 119 | modelG.module.update_training_batch((start_epoch-1)//opt.niter_step) |
| 120 | |
| 121 | n_gpus = opt.n_gpus_gen if opt.batchSize == 1 else 1 # number of gpus used for generator for each batch |
| 122 | tG, tD = opt.n_frames_G, opt.n_frames_D |
| 123 | tDB = tD * opt.output_nc |
| 124 | s_scales = opt.n_scales_spatial |
| 125 | t_scales = opt.n_scales_temporal |
| 126 | input_nc = 1 if opt.label_nc != 0 else opt.input_nc |
| 127 | output_nc = opt.output_nc |
| 128 | |
| 129 | print_freq = lcm(opt.print_freq, opt.batchSize) |
| 130 | total_steps = (start_epoch-1) * len(data_loader) + epoch_iter |
| 131 | total_steps = total_steps // print_freq * print_freq |
| 132 | |
| 133 | return n_gpus, tG, tD, tDB, s_scales, t_scales, input_nc, output_nc, start_epoch, epoch_iter, print_freq, total_steps, iter_path |
| 134 | |
| 135 | def save_models(opt, epoch, epoch_iter, total_steps, visualizer, iter_path, modelG, modelD, end_of_epoch=False): |
| 136 | if not end_of_epoch: |
no test coverage detected