Team Ai
Apppublic

Arulkumar03/Fox_Sheep_Detector_Computer_Vision_model

sourceHugging Facemitupdated 3y agoView on Hugging Face
0likes
train_net.py81 linesDownload Raw Back to DensePose
1#!/usr/bin/env python32# Copyright (c) Facebook, Inc. and its affiliates.3 4"""5DensePose Training Script.6 7This script is similar to the training script in detectron2/tools.8 9It is an example of how a user might use detectron2 for a new project.10"""11 12from datetime import timedelta13 14import detectron2.utils.comm as comm15from detectron2.config import get_cfg16from detectron2.engine import DEFAULT_TIMEOUT, default_argument_parser, default_setup, hooks, launch17from detectron2.evaluation import verify_results18from detectron2.utils.file_io import PathManager19from detectron2.utils.logger import setup_logger20 21from densepose import add_densepose_config22from densepose.engine import Trainer23from densepose.modeling.densepose_checkpoint import DensePoseCheckpointer24 25 26def setup(args):27    cfg = get_cfg()28    add_densepose_config(cfg)29    cfg.merge_from_file(args.config_file)30    cfg.merge_from_list(args.opts)31    cfg.freeze()32    default_setup(cfg, args)33    # Setup logger for "densepose" module34    setup_logger(output=cfg.OUTPUT_DIR, distributed_rank=comm.get_rank(), name="densepose")35    return cfg36 37 38def main(args):39    cfg = setup(args)40    # disable strict kwargs checking: allow one to specify path handle41    # hints through kwargs, like timeout in DP evaluation42    PathManager.set_strict_kwargs_checking(False)43 44    if args.eval_only:45        model = Trainer.build_model(cfg)46        DensePoseCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(47            cfg.MODEL.WEIGHTS, resume=args.resume48        )49        res = Trainer.test(cfg, model)50        if cfg.TEST.AUG.ENABLED:51            res.update(Trainer.test_with_TTA(cfg, model))52        if comm.is_main_process():53            verify_results(cfg, res)54        return res55 56    trainer = Trainer(cfg)57    trainer.resume_or_load(resume=args.resume)58    if cfg.TEST.AUG.ENABLED:59        trainer.register_hooks(60            [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]61        )62    return trainer.train()63 64 65if __name__ == "__main__":66    args = default_argument_parser().parse_args()67    cfg = setup(args)68    timeout = (69        DEFAULT_TIMEOUT if cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE else timedelta(hours=4)70    )71    print("Command Line Args:", args)72    launch(73        main,74        args.num_gpus,75        num_machines=args.num_machines,76        machine_rank=args.machine_rank,77        dist_url=args.dist_url,78        args=(args,),79        timeout=timeout,80    )81