File size: 8,994 Bytes
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
c74a070
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
a80d6bb
 
c74a070
 
a80d6bb
 
c74a070
 
a80d6bb
 
c74a070
a80d6bb
 
 
 
c74a070
a80d6bb
 
 
 
 
 
 
 
c74a070
 
 
 
 
 
 
 
 
a80d6bb
 
 
c74a070
 
 
a80d6bb
 
c74a070
 
 
a80d6bb
 
c74a070
 
 
 
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
 
 
 
 
 
 
 
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
c74a070
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a80d6bb
 
 
 
 
 
 
 
 
 
 
c74a070
 
 
 
 
a80d6bb
c74a070
 
a80d6bb
 
 
c74a070
a80d6bb
 
 
 
 
 
c74a070
a80d6bb
c74a070
a80d6bb
c74a070
a80d6bb
 
 
 
 
 
 
 
 
 
 
 
 
 
c74a070
a80d6bb
 
 
c74a070
a80d6bb
 
 
c74a070
a80d6bb
c74a070
a80d6bb
 
 
c74a070
a80d6bb
 
 
 
 
c74a070
a80d6bb
 
 
 
c74a070
 
 
 
 
 
 
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
"""
Main file to launch training and testing experiments.
"""

import yaml
import os
import argparse
import numpy as np
import torch

from .config.project_config import Config as cfg
from .train import train_net
from .export import export_predictions, export_homograpy_adaptation


# Pytorch configurations
torch.cuda.empty_cache()
torch.backends.cudnn.benchmark = True


def load_config(config_path):
    """Load configurations from a given yaml file."""
    # Check file exists
    if not os.path.exists(config_path):
        raise ValueError("[Error] The provided config path is not valid.")

    # Load the configuration
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    return config


def update_config(path, model_cfg=None, dataset_cfg=None):
    """Update configuration file from the resume path."""
    # Check we need to update or completely override.
    model_cfg = {} if model_cfg is None else model_cfg
    dataset_cfg = {} if dataset_cfg is None else dataset_cfg

    # Load saved configs
    with open(os.path.join(path, "model_cfg.yaml"), "r") as f:
        model_cfg_saved = yaml.safe_load(f)
        model_cfg.update(model_cfg_saved)
    with open(os.path.join(path, "dataset_cfg.yaml"), "r") as f:
        dataset_cfg_saved = yaml.safe_load(f)
        dataset_cfg.update(dataset_cfg_saved)

    # Update the saved yaml file
    if not model_cfg == model_cfg_saved:
        with open(os.path.join(path, "model_cfg.yaml"), "w") as f:
            yaml.dump(model_cfg, f)
    if not dataset_cfg == dataset_cfg_saved:
        with open(os.path.join(path, "dataset_cfg.yaml"), "w") as f:
            yaml.dump(dataset_cfg, f)

    return model_cfg, dataset_cfg


def record_config(model_cfg, dataset_cfg, output_path):
    """Record dataset config to the log path."""
    # Record model config
    with open(os.path.join(output_path, "model_cfg.yaml"), "w") as f:
        yaml.safe_dump(model_cfg, f)

    # Record dataset config
    with open(os.path.join(output_path, "dataset_cfg.yaml"), "w") as f:
        yaml.safe_dump(dataset_cfg, f)


def train(args, dataset_cfg, model_cfg, output_path):
    """Training function."""
    # Update model config from the resume path (only in resume mode)
    if args.resume:
        if os.path.realpath(output_path) != os.path.realpath(args.resume_path):
            record_config(model_cfg, dataset_cfg, output_path)

    # First time, then write the config file to the output path
    else:
        record_config(model_cfg, dataset_cfg, output_path)

    # Launch the training
    train_net(args, dataset_cfg, model_cfg, output_path)


def export(
    args,
    dataset_cfg,
    model_cfg,
    output_path,
    export_dataset_mode=None,
    device=torch.device("cuda"),
):
    """Export function."""
    # Choose between normal predictions export or homography adaptation
    if dataset_cfg.get("homography_adaptation") is not None:
        print("[Info] Export predictions with homography adaptation.")
        export_homograpy_adaptation(
            args, dataset_cfg, model_cfg, output_path, export_dataset_mode, device
        )
    else:
        print("[Info] Export predictions normally.")
        export_predictions(
            args, dataset_cfg, model_cfg, output_path, export_dataset_mode
        )


def main(
    args, dataset_cfg, model_cfg, export_dataset_mode=None, device=torch.device("cuda")
):
    """Main function."""
    # Make the output path
    output_path = os.path.join(cfg.EXP_PATH, args.exp_name)

    if args.mode == "train":
        if not os.path.exists(output_path):
            os.makedirs(output_path)
        print("[Info] Training mode")
        print("\t Output path: %s" % output_path)
        train(args, dataset_cfg, model_cfg, output_path)
    elif args.mode == "export":
        # Different output_path in export mode
        output_path = os.path.join(cfg.export_dataroot, args.exp_name)
        print("[Info] Export mode")
        print("\t Output path: %s" % output_path)
        export(
            args,
            dataset_cfg,
            model_cfg,
            output_path,
            export_dataset_mode,
            device=device,
        )
    else:
        raise ValueError("[Error]: Unknown mode: " + args.mode)


def set_random_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)


if __name__ == "__main__":
    # Parse input arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--mode", type=str, default="train", help="'train' or 'export'."
    )
    parser.add_argument(
        "--dataset_config", type=str, default=None, help="Path to the dataset config."
    )
    parser.add_argument(
        "--model_config", type=str, default=None, help="Path to the model config."
    )
    parser.add_argument("--exp_name", type=str, default="exp", help="Experiment name.")
    parser.add_argument(
        "--resume",
        action="store_true",
        default=False,
        help="Load a previously trained model.",
    )
    parser.add_argument(
        "--pretrained",
        action="store_true",
        default=False,
        help="Start training from a pre-trained model.",
    )
    parser.add_argument(
        "--resume_path", default=None, help="Path from which to resume training."
    )
    parser.add_argument(
        "--pretrained_path", default=None, help="Path to the pre-trained model."
    )
    parser.add_argument(
        "--checkpoint_name", default=None, help="Name of the checkpoint to use."
    )
    parser.add_argument(
        "--export_dataset_mode", default=None, help="'train' or 'test'."
    )
    parser.add_argument(
        "--export_batch_size", default=4, type=int, help="Export batch size."
    )

    args = parser.parse_args()

    # Check if GPU is available
    # Get the model
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    # Check if dataset config and model config is given.
    if (
        ((args.dataset_config is None) or (args.model_config is None))
        and (not args.resume)
        and (args.mode == "train")
    ):
        raise ValueError(
            "[Error] The dataset config and model config should be given in non-resume mode"
        )

    # If resume, check if the resume path has been given
    if args.resume and (args.resume_path is None):
        raise ValueError("[Error] Missing resume path.")

    # [Training] Load the config file.
    if args.mode == "train" and (not args.resume):
        # Check the pretrained checkpoint_path exists
        if args.pretrained:
            checkpoint_folder = args.resume_path
            checkpoint_path = os.path.join(args.pretrained_path, args.checkpoint_name)
            if not os.path.exists(checkpoint_path):
                raise ValueError("[Error] Missing checkpoint: " + checkpoint_path)
        dataset_cfg = load_config(args.dataset_config)
        model_cfg = load_config(args.model_config)

    # [resume Training, Test, Export] Load the config file.
    elif (args.mode == "train" and args.resume) or (args.mode == "export"):
        # Check checkpoint path exists
        checkpoint_folder = args.resume_path
        checkpoint_path = os.path.join(args.resume_path, args.checkpoint_name)
        if not os.path.exists(checkpoint_path):
            raise ValueError("[Error] Missing checkpoint: " + checkpoint_path)

        # Load model_cfg from checkpoint folder if not provided
        if args.model_config is None:
            print("[Info] No model config provided. Loading from checkpoint folder.")
            model_cfg_path = os.path.join(checkpoint_folder, "model_cfg.yaml")
            if not os.path.exists(model_cfg_path):
                raise ValueError("[Error] Missing model config in checkpoint path.")
            model_cfg = load_config(model_cfg_path)
        else:
            model_cfg = load_config(args.model_config)

        # Load dataset_cfg from checkpoint folder if not provided
        if args.dataset_config is None:
            print("[Info] No dataset config provided. Loading from checkpoint folder.")
            dataset_cfg_path = os.path.join(checkpoint_folder, "dataset_cfg.yaml")
            if not os.path.exists(dataset_cfg_path):
                raise ValueError("[Error] Missing dataset config in checkpoint path.")
            dataset_cfg = load_config(dataset_cfg_path)
        else:
            dataset_cfg = load_config(args.dataset_config)

        # Check the --export_dataset_mode flag
        if (args.mode == "export") and (args.export_dataset_mode is None):
            raise ValueError("[Error] Empty --export_dataset_mode flag.")
    else:
        raise ValueError("[Error] Unknown mode: " + args.mode)

    # Set the random seed
    seed = dataset_cfg.get("random_seed", 0)
    set_random_seed(seed)

    main(
        args,
        dataset_cfg,
        model_cfg,
        export_dataset_mode=args.export_dataset_mode,
        device=device,
    )