train_ds, epochs=cfg.train.epochs, batch_size=cfg.train.batch_size, verbose=cfg.train.verbose, validation_data=None if not cfg.train.val else val_ds, validation_steps=val_steps[cfg.data.dataset] // strategy.num_replicas_in_sync, steps_per_epoch=spe) yolo.save_weights(weights_path='./models/{}/weights'.format( cfg.model.name), weights_type=cfg.train.save_weights_type) pickle.dump(hist.history, open('./models/{}/history.pkl'.format(cfg.model.name), 'wb')) return yolo if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-c', '--cfg', default='default') args = parser.parse_args() cfg = CN(new_allowed=True) cfg.merge_from_file(osp.join('configs', args.cfg + '.yaml')) cfg.model.name = args.cfg tpu, strategy = detect_hardware(tpu_name=None) yolo = train(cfg, strategy) predict(yolo, cfg, dataset=cfg.data.dataset, split='val')
import numpy as np from lr_schedules import WarmupCosineDecay, WarmupPiecewise import os.path as osp from utils import detect_hardware from dataset.dataloader import load_representative_tfds, visualize, prediction_tf_lite, prediction_examples from dataset.coco import cn as cfg import argparse from lr_schedules import WarmupCosineDecay import math if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-c', '--cfg', required=True) args = parser.parse_args() tpu, strategy = detect_hardware(None) if tpu: cfg.TRAIN.ACCELERATOR = args.tpu else: cfg.TRAIN.ACCELERATOR = 'GPU/CPU' cfg.merge_from_file('configs/' + args.cfg) cfg.MODEL.NAME = args.cfg.split('.yaml')[0] if cfg.DATASET.OUTPUT_SHAPE[-1] == 17: cfg.DATASET.KP_FLIP = cfg.DATASET.KP_FLIP[:17] cfg.MODEL.TFRECORDS = 'data/tfrecords_foot/val' train_ds = load_representative_tfds(cfg) def representative_dataset(): for img in train_ds:
AP_large = cocoEval.stats[5] return mAP, AP_50, AP_75, AP_small, AP_medium, AP_large # AP if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--tpu', default=None) parser.add_argument('-c', '--cfg', required=True) # yaml parser.add_argument('--det', type=int, default=-1) parser.add_argument('--ckpt', default=None) parser.add_argument('--split', default='val') args = parser.parse_args() from dataset.coco import cn as cfg cfg.merge_from_file('configs/' + args.cfg) cfg.MODEL.NAME = args.cfg.split('.')[0] if cfg.DATASET.OUTPUT_SHAPE[-1] == 17: cfg.DATASET.KP_FLIP = cfg.DATASET.KP_FLIP[:17] if args.ckpt: cfg.MODEL.NAME += '_{}'.format(args.ckpt) if args.det >= 0: cfg.VAL.DET = bool(args.det) tpu, strategy = detect_hardware(args.tpu) if args.split == 'val': AP, _ = validate(strategy, cfg, split='val') print('AP: {:.5f}'.format(AP)) else: validate(strategy, cfg, split='test')