Spaces:
Sleeping
Sleeping
| import torch | |
| from torch.utils.data import DataLoader | |
| import os | |
| from train.trainer_step import TrainStepper | |
| from train.base_trainer import trainer, evaluator | |
| from data.base_dataset import BaseDataset | |
| from data.mixed_dataset import MixedDataset | |
| from models.deco import DECO | |
| from utils.config import parse_args, run_grid_search_experiments | |
| def train(hparams): | |
| deco_model = DECO(hparams.TRAINING.ENCODER, hparams.TRAINING.CONTEXT, device) | |
| solver = TrainStepper(deco_model, hparams.TRAINING.CONTEXT, hparams.OPTIMIZER.LR, hparams.TRAINING.LOSS_WEIGHTS, hparams.TRAINING.PAL_LOSS_WEIGHTS, device) | |
| vb_f1 = 0 | |
| start_ep = 0 | |
| num = 0 | |
| k = True | |
| latest_model_path = hparams.TRAINING.BEST_MODEL_PATH.replace('best', 'latest') | |
| if os.path.exists(latest_model_path): | |
| _, vb_f1 = solver.load(hparams.TRAINING.BEST_MODEL_PATH) | |
| start_ep, _ = solver.load(latest_model_path) | |
| for epoch in range(start_ep+1, hparams.TRAINING.NUM_EPOCHS + 1): | |
| # Train one epoch | |
| trainer(epoch, train_loader, solver, hparams) | |
| # Run evaluation | |
| vc_f1 = None | |
| for val_loader in val_loaders: | |
| dataset_name = val_loader.dataset.dataset | |
| vc_f1_ds = evaluator(val_loader, solver, hparams, epoch, dataset_name, normalize=hparams.DATASET.NORMALIZE_IMAGES) | |
| if dataset_name == hparams.VALIDATION.MAIN_DATASET: | |
| vc_f1 = vc_f1_ds | |
| if vc_f1 is None: | |
| raise ValueError('Main dataset not found in validation datasets') | |
| print('Learning rate: ', solver.lr) | |
| print('---------------------------------------------') | |
| print('---------------------------------------------') | |
| solver.save(epoch, vc_f1, latest_model_path) | |
| if epoch % hparams.TRAINING.CHECKPOINT_EPOCHS == 0: | |
| inter_model_path = latest_model_path.replace('latest', 'epoch_'+str(epoch).zfill(3)) | |
| solver.save(epoch, vc_f1, inter_model_path) | |
| if vc_f1 < vb_f1: | |
| num += 1 | |
| print('Not Saving model: Best Val F1 = ', vb_f1, ' Current Val F1 = ', vc_f1) | |
| else: | |
| num = 0 | |
| vb_f1 = vc_f1 | |
| print('Saving model...') | |
| solver.save(epoch, vb_f1, hparams.TRAINING.BEST_MODEL_PATH) | |
| if num >= hparams.OPTIMIZER.NUM_UPDATE_LR: solver.update_lr() | |
| if num >= hparams.TRAINING.NUM_EARLY_STOP: | |
| print('Early Stop') | |
| k = False | |
| if k: continue | |
| else: break | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| hparams = run_grid_search_experiments( | |
| args, | |
| script='train.py', | |
| ) | |
| if torch.cuda.is_available(): | |
| device = torch.device('cuda') | |
| else: | |
| device = torch.device('cpu') | |
| train_dataset = MixedDataset(hparams.TRAINING.DATASETS, 'train', dataset_mix_pdf=hparams.TRAINING.DATASET_MIX_PDF, normalize=hparams.DATASET.NORMALIZE_IMAGES) | |
| val_datasets = [] | |
| for ds in hparams.VALIDATION.DATASETS: | |
| if ds in ['rich', 'prox']: | |
| val_datasets.append(BaseDataset(ds, 'val', model_type='smplx', normalize=hparams.DATASET.NORMALIZE_IMAGES)) | |
| elif ds in ['damon']: | |
| val_datasets.append(BaseDataset(ds, 'val', model_type='smpl', normalize=hparams.DATASET.NORMALIZE_IMAGES)) | |
| else: | |
| raise ValueError('Dataset not supported') | |
| train_loader = DataLoader(train_dataset, hparams.DATASET.BATCH_SIZE, shuffle=True, num_workers=hparams.DATASET.NUM_WORKERS) | |
| val_loaders = [DataLoader(val_dataset, batch_size=hparams.DATASET.BATCH_SIZE, shuffle=False, num_workers=hparams.DATASET.NUM_WORKERS) for val_dataset in val_datasets] | |
| train(hparams) | |