| import os |
| import sys |
| import pathlib |
|
|
| __dir__ = pathlib.Path(os.path.abspath(__file__)) |
| sys.path.append(str(__dir__)) |
| sys.path.append(str(__dir__.parent.parent)) |
|
|
| import paddle |
| import paddle.distributed as dist |
| from utils import Config, ArgsParser |
|
|
|
|
| def init_args(): |
| parser = ArgsParser() |
| args = parser.parse_args() |
| return args |
|
|
|
|
| def main(config, profiler_options): |
| from models import build_model, build_loss |
| from data_loader import get_dataloader |
| from trainer import Trainer |
| from post_processing import get_post_processing |
| from utils import get_metric |
|
|
| if paddle.device.cuda.device_count() > 1: |
| dist.init_parallel_env() |
| config["distributed"] = True |
| else: |
| config["distributed"] = False |
| train_loader = get_dataloader(config["dataset"]["train"], config["distributed"]) |
| assert train_loader is not None |
| if "validate" in config["dataset"]: |
| validate_loader = get_dataloader(config["dataset"]["validate"], False) |
| else: |
| validate_loader = None |
| criterion = build_loss(config["loss"]) |
| config["arch"]["backbone"]["in_channels"] = ( |
| 3 if config["dataset"]["train"]["dataset"]["args"]["img_mode"] != "GRAY" else 1 |
| ) |
| model = build_model(config["arch"]) |
| |
| post_p = get_post_processing(config["post_processing"]) |
| metric = get_metric(config["metric"]) |
| trainer = Trainer( |
| config=config, |
| model=model, |
| criterion=criterion, |
| train_loader=train_loader, |
| post_process=post_p, |
| metric_cls=metric, |
| validate_loader=validate_loader, |
| profiler_options=profiler_options, |
| ) |
| trainer.train() |
|
|
|
|
| if __name__ == "__main__": |
| args = init_args() |
| assert os.path.exists(args.config_file) |
| config = Config(args.config_file) |
| config.merge_dict(args.opt) |
| main(config.cfg, args.profiler_options) |
|
|
