Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,386 Bytes
184241a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import re
import subprocess
import os
import argparse
from mmcv import Config, DictAction
def init_parser():
# Get config and work_dir from user
parser = argparse.ArgumentParser(description='Run the pipeline')
parser.add_argument('--config', help='config file', required=True)
parser.add_argument('--work_dir', help='work directory', required=True)
parser.add_argument('--best', action='store_true', help='work directory')
parser.add_argument('--supervision', type=str, default='decoder', help='adj supervision')
parser.add_argument('--ft_epochs', type=int, default=100, help='work directory')
parser.add_argument('--masking_ratio', type=float, default=0.5, help='work directory')
parser.add_argument('--lamda_masking', type=float, default=1.0, help='work directory')
args = parser.parse_args()
return args
def get_best_model(work_dir):
if os.path.exists(work_dir):
file_names = [filename for filename in os.listdir(work_dir) if filename.startswith("best_")]
if len(file_names) > 0:
file_name = file_names[0]
ckpt_path = f'{work_dir}/{file_name}'
else:
ckpt_path = f'{work_dir}/latest.pth'
return ckpt_path
def main():
args = init_parser()
config = args.config
work_dir = args.work_dir
if args.best:
work_dir = f'{work_dir}_best_ckpt'
if not os.path.exists(work_dir):
os.makedirs(work_dir)
subprocess.run(['cp', config, work_dir])
# -----------------------------BASE MODEL TRAINING--------------------------------
base_workdir = f'{work_dir}/base'
cfg = Config.fromfile(args.config)
num_epochs = cfg.total_epochs
final_epoch_path = f'{base_workdir}/epoch_{num_epochs}.pth'
if not os.path.exists(final_epoch_path):
print("Running Base Model Training")
subprocess.run(['python', 'train.py', '--config', config, '--work-dir', base_workdir])
# -----------------------------SKELETON MODEL TRAINING--------------------------------
skeleton_work_dir = f'{work_dir}/base_skeleton'
skeleton_final_epoch_path = f'{skeleton_work_dir}/epoch_{args.ft_epochs}.pth'
if args.best:
best_ckpt = get_best_model(base_workdir)
load_from = best_ckpt
else:
load_from = final_epoch_path
new_cfg = Config.fromfile(args.config)
new_cfg.load_from = load_from
new_cfg.total_epochs = args.ft_epochs
new_cfg.model.freeze_backbone = True
new_cfg.model.keypoint_head.skeleton_head['learn_skeleton'] = True
new_cfg.model.keypoint_head.learn_skeleton = True
new_cfg.model.keypoint_head.masked_supervision = True
new_cfg.model.keypoint_head.masking_ratio = args.masking_ratio
new_cfg.model.keypoint_head.skeleton_loss_weight = args.lamda_masking
Config.dump(new_cfg, f'{work_dir}/skeleton_config.py')
if not os.path.exists(skeleton_final_epoch_path):
print("Running Base Model Training")
subprocess.run(
['python', 'train.py', '--config', f'{work_dir}/skeleton_config.py', '--work-dir', skeleton_work_dir])
# -----------------------------BIAS MODEL TRAINING--------------------------------
bias_work_dir = f'{work_dir}/base_skeleton_bias'
bias_final_epoch_path = f'{bias_work_dir}/epoch_{args.ft_epochs}.pth'
if args.best:
best_ckpt = get_best_model(skeleton_work_dir)
load_from = best_ckpt
else:
load_from = skeleton_final_epoch_path
new_cfg.load_from = load_from
new_cfg.model.keypoint_head.transformer.use_bias_attn_module = True
new_cfg.model.keypoint_head.transformer.attn_bias = True
new_cfg.model.keypoint_head.transformer.max_hops = 4
new_cfg.model.keypoint_head.model_freeze = 'skeleton'
Config.dump(new_cfg, f'{work_dir}/bias_config.py')
if not os.path.exists(bias_final_epoch_path):
print("Running Bias Model Training")
subprocess.run(
['python', 'train.py', '--config', f'{work_dir}/bias_config.py', '--work-dir', bias_work_dir])
# -----------------------------EVALUATION--------------------------------
best_ckpt = get_best_model(bias_work_dir)
subprocess.run(['python', 'test.py', f'{work_dir}/bias_config.py', f'{bias_work_dir}/latest.pth'])
subprocess.run(['python', 'test.py', f'{work_dir}/bias_config.py', best_ckpt])
if __name__ == '__main__':
main()
|