-
Notifications
You must be signed in to change notification settings - Fork 7
/
train.py
69 lines (51 loc) · 2.29 KB
/
train.py
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
"""Anomalib Traning Script.
This script reads the name of the model or config file from command
line, train/test the anomaly model to get quantitative and qualitative
results.
"""
# Copyright (C) 2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import logging
import warnings
from argparse import ArgumentParser, Namespace
from pytorch_lightning import Trainer, seed_everything
from anomalib.config import get_configurable_parameters
from anomalib.data import get_datamodule
from anomalib.models import get_model
from anomalib.utils.callbacks import LoadModelCallback, get_callbacks
from anomalib.utils.loggers import configure_logger, get_experiment_logger
logger = logging.getLogger("anomalib")
def get_args() -> Namespace:
"""Get command line arguments.
Returns:
Namespace: List of arguments.
"""
parser = ArgumentParser()
parser.add_argument("--model", type=str, default="patchcore", help="Name of the algorithm to train/test")
parser.add_argument("--config", type=str, default="configs/aqa.yaml", help="Path to a model config file")
parser.add_argument("--log-level", type=str, default="INFO", help="<DEBUG, INFO, WARNING, ERROR>")
args = parser.parse_args()
return args
def train():
"""Train an anomaly classification or segmentation model based on a provided configuration file."""
args = get_args()
configure_logger(level=args.log_level)
if args.log_level == "ERROR":
warnings.filterwarnings("ignore")
config = get_configurable_parameters(model_name=args.model, config_path=args.config)
if config.project.seed:
seed_everything(config.project.seed)
datamodule = get_datamodule(config)
model = get_model(config)
experiment_logger = get_experiment_logger(config)
callbacks = get_callbacks(config)
trainer = Trainer(**config.trainer, logger=experiment_logger, callbacks=callbacks)
logger.info("Training the model.")
trainer.fit(model=model, datamodule=datamodule)
logger.info("Loading the best model weights.")
load_model_callback = LoadModelCallback(weights_path=trainer.checkpoint_callback.best_model_path)
trainer.callbacks.insert(0, load_model_callback)
logger.info("Testing the model.")
trainer.test(model=model, datamodule=datamodule)
if __name__ == "__main__":
train()