ChangeDETR/main.py at main · hustvl/ChangeDETR · GitHub

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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
import warnings
from argparse import ArgumentParser

import pytorch_lightning as pl
import torch
from easydict import EasyDict
from loguru import logger as L
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.plugins import DDPPlugin
from pytorch_lightning.utilities import rank_zero_only

from data.datamodule import DataModule
from models.centernet_with_coam import CenterNetWithCoAttention
from models.defdetr import DefDetr
from utils.general import get_easy_dict_from_yaml_file

warnings.filterwarnings("ignore")


@rank_zero_only
def print_args(configs):
    L.log("INFO", configs)


def train(configs, model, logger, datamodule, callbacks=None):
    L.log("INFO", f"Training model.")
    trainer = pl.Trainer.from_argparse_args(
        configs,
        logger=logger,
        strategy=DDPPlugin(find_unused_parameters=False),
        log_every_n_steps=1,
        callbacks=callbacks,
        check_val_every_n_epoch=1,
        benchmark=False,
        profiler="simple",
    )
    trainer.fit(model, datamodule=datamodule)
    return trainer, trainer.checkpoint_callback.best_model_path


def test(configs, model, logger, datamodule, checkpoint_path, callbacks=None):
    L.log("INFO", f"Testing model.")
    tester = pl.Trainer.from_argparse_args(
        configs, logger=logger, callbacks=callbacks, benchmark=True
    )
    tester.test(model=model, datamodule=datamodule, ckpt_path=checkpoint_path)

def get_logging_callback_manager(args):
    if args.method == "centernet":
        from models.centernet_with_coam import TensorboardCallbackManager
        return TensorboardCallbackManager(args)
    elif args.method == 'defdetr':
        from models.defdetr import TensorboardCallbackManager
        return TensorboardCallbackManager(args)
    else:
        raise NotImplementedError(f"Given method ({args.method}) not implemented!")

if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--method", type=str)
    parser.add_argument("--no_logging", action="store_true", default=False)
    parser.add_argument("--test_from_checkpoint", type=str, default="")
    parser.add_argument("--quick_prototype", action="store_true", default=False)
    parser.add_argument("--pretrained_weights", type=str, default=None)
    parser.add_argument("--load_weights_from", type=str, default=None)
    parser.add_argument("--config_file", required=True)
    parser.add_argument("--experiment_name", type=str, default=None)
    args, _ = parser.parse_known_args()
    if args.method == "centernet":
        parser = CenterNetWithCoAttention.add_model_specific_args(parser)
    elif args.method == 'defdetr':
        parser = DefDetr.add_model_specific_args(parser) # 类方法，未初始化实例就可以使用
    else:
        raise NotImplementedError(f"Unknown method type {args.method}")

    # parse configs from cmd line and config file into an EasyDict
    parser = DataModule.add_data_specific_args(parser)
    parser = pl.Trainer.add_argparse_args(parser)
    args = parser.parse_args()
    args = EasyDict(vars(args))
    configs = get_easy_dict_from_yaml_file(args.config_file)

    # copy cmd line configs into config file configs, overwriting any duplicates
    for key in args.keys():
        if args[key] is not None:
            configs[key] = args[key]
        elif key not in configs.keys():
            configs[key] = args[key]

    if configs.quick_prototype:
        configs.limit_train_batches = 200
        configs.limit_val_batches = 20
        configs.limit_test_batches = 20
        configs.max_epochs = 20

    print_args(configs)

    pl.seed_everything(42, workers=True)

    datamodule = DataModule(configs)
    if configs.method == "centernet":
        model = CenterNetWithCoAttention(configs)
    elif configs.method == 'defdetr':
        model = DefDetr(configs)

    logger = None
    from pytorch_lightning.callbacks import LearningRateMonitor
    lr_monitor = LearningRateMonitor(logging_interval='epoch')
    callbacks = [get_logging_callback_manager(configs), lr_monitor]
    if not configs.no_logging:
        logger = TensorBoardLogger(
            name=configs.experiment_name,
            save_dir=configs.log_dir,
        )
        callbacks.append(ModelCheckpoint(dirpath=configs.get('ckpt_dir', None),
                                         filename="epoch{epoch}-valloss{val/overall_loss:.2f}",
                                         auto_insert_metric_name=False,
                                         monitor="val/overall_loss",
                                         mode="min",
                                         save_last=True))

    trainer = None
    if configs.test_from_checkpoint == "":
        # train the model and store the path to the best model (as per the validation set)
        # Note: multi-GPU training is supported.
        trainer, test_checkpoint_path = train(configs, model, logger, datamodule, callbacks)
        # test the best model exactly once on a single GPU
        torch.distributed.destroy_process_group()
    else:
        # test the given model checkpoint
        test_checkpoint_path = configs.test_from_checkpoint

    configs.gpus = 1
    if trainer is None or trainer.global_rank == 0:
        test(
            configs,
            model,
            logger,
            datamodule,
            test_checkpoint_path if test_checkpoint_path != "" else None,
            callbacks,
        )