INNER CODE UNIT · Python

main

FoundationVision/UniTok · main.py:170

def main():
    args = config.init_dist_and_get_args()
    print(f'[args] initial args:\n{str(args)}')

    # resume ckpt
    ckpt = maybe_auto_resume(args, 'ckpt*.pth')
    start_iter = ckpt.get('iter', 0)
    start_epoch = ckpt.get('epoch', 0)
    trainer_state = ckpt.get('trainer', {})

    # load data
    print(f'[data] load data...\n')
    aug_cfg = {'scale': [0.64, 1.0]} if args.use_aug else None
    preprocess_fns = build_clip_transforms(args, aug_cfg=aug_cfg)
    tokenizer = partial(tokenize, context_length=args.text_context_length)
    data = load_data(args, preprocess_fns, epoch=start_epoch, iters=start_iter, tokenizer=tokenizer)

    # build models

View source record →