implemented inference
This commit is contained in:
Jan Kowalczyk
@@ -3,6 +3,7 @@ import torch
|
||||
import logging
|
||||
import random
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
from utils.config import Config
|
||||
from utils.visualization.plot_images_grid import plot_images_grid
|
||||
@@ -14,6 +15,15 @@ from datasets.main import load_dataset
|
||||
# Settings
|
||||
################################################################################
|
||||
@click.command()
|
||||
@click.argument(
|
||||
"action",
|
||||
type=click.Choice(
|
||||
[
|
||||
"train",
|
||||
"infer",
|
||||
]
|
||||
),
|
||||
)
|
||||
@click.argument(
|
||||
"dataset_name",
|
||||
type=click.Choice(
|
||||
@@ -203,6 +213,7 @@ from datasets.main import load_dataset
|
||||
"If > 1, the specified number of outlier classes will be sampled at random.",
|
||||
)
|
||||
def main(
|
||||
action,
|
||||
dataset_name,
|
||||
net_name,
|
||||
xp_path,
|
||||
@@ -303,138 +314,194 @@ def main(
|
||||
logger.info("Number of threads: %d" % num_threads)
|
||||
logger.info("Number of dataloader workers: %d" % n_jobs_dataloader)
|
||||
|
||||
# Load data
|
||||
dataset = load_dataset(
|
||||
dataset_name,
|
||||
data_path,
|
||||
normal_class,
|
||||
known_outlier_class,
|
||||
n_known_outlier_classes,
|
||||
ratio_known_normal,
|
||||
ratio_known_outlier,
|
||||
ratio_pollution,
|
||||
random_state=np.random.RandomState(cfg.settings["seed"]),
|
||||
)
|
||||
# Log random sample of known anomaly classes if more than 1 class
|
||||
if n_known_outlier_classes > 1:
|
||||
logger.info("Known anomaly classes: %s" % (dataset.known_outlier_classes,))
|
||||
if action == "train":
|
||||
|
||||
# Initialize DeepSAD model and set neural network phi
|
||||
deepSAD = DeepSAD(cfg.settings["eta"])
|
||||
deepSAD.set_network(net_name)
|
||||
|
||||
# If specified, load Deep SAD model (center c, network weights, and possibly autoencoder weights)
|
||||
if load_model:
|
||||
deepSAD.load_model(model_path=load_model, load_ae=True, map_location=device)
|
||||
logger.info("Loading model from %s." % load_model)
|
||||
|
||||
logger.info("Pretraining: %s" % pretrain)
|
||||
if pretrain:
|
||||
# Log pretraining details
|
||||
logger.info("Pretraining optimizer: %s" % cfg.settings["ae_optimizer_name"])
|
||||
logger.info("Pretraining learning rate: %g" % cfg.settings["ae_lr"])
|
||||
logger.info("Pretraining epochs: %d" % cfg.settings["ae_n_epochs"])
|
||||
logger.info(
|
||||
"Pretraining learning rate scheduler milestones: %s"
|
||||
% (cfg.settings["ae_lr_milestone"],)
|
||||
# Load data
|
||||
dataset = load_dataset(
|
||||
dataset_name,
|
||||
data_path,
|
||||
normal_class,
|
||||
known_outlier_class,
|
||||
n_known_outlier_classes,
|
||||
ratio_known_normal,
|
||||
ratio_known_outlier,
|
||||
ratio_pollution,
|
||||
random_state=np.random.RandomState(cfg.settings["seed"]),
|
||||
)
|
||||
logger.info("Pretraining batch size: %d" % cfg.settings["ae_batch_size"])
|
||||
logger.info("Pretraining weight decay: %g" % cfg.settings["ae_weight_decay"])
|
||||
# Log random sample of known anomaly classes if more than 1 class
|
||||
if n_known_outlier_classes > 1:
|
||||
logger.info("Known anomaly classes: %s" % (dataset.known_outlier_classes,))
|
||||
|
||||
# Pretrain model on dataset (via autoencoder)
|
||||
deepSAD.pretrain(
|
||||
# Initialize DeepSAD model and set neural network phi
|
||||
deepSAD = DeepSAD(cfg.settings["eta"])
|
||||
deepSAD.set_network(net_name)
|
||||
|
||||
# If specified, load Deep SAD model (center c, network weights, and possibly autoencoder weights)
|
||||
if load_model:
|
||||
deepSAD.load_model(model_path=load_model, load_ae=True, map_location=device)
|
||||
logger.info("Loading model from %s." % load_model)
|
||||
|
||||
logger.info("Pretraining: %s" % pretrain)
|
||||
if pretrain:
|
||||
# Log pretraining details
|
||||
logger.info("Pretraining optimizer: %s" % cfg.settings["ae_optimizer_name"])
|
||||
logger.info("Pretraining learning rate: %g" % cfg.settings["ae_lr"])
|
||||
logger.info("Pretraining epochs: %d" % cfg.settings["ae_n_epochs"])
|
||||
logger.info(
|
||||
"Pretraining learning rate scheduler milestones: %s"
|
||||
% (cfg.settings["ae_lr_milestone"],)
|
||||
)
|
||||
logger.info("Pretraining batch size: %d" % cfg.settings["ae_batch_size"])
|
||||
logger.info(
|
||||
"Pretraining weight decay: %g" % cfg.settings["ae_weight_decay"]
|
||||
)
|
||||
|
||||
# Pretrain model on dataset (via autoencoder)
|
||||
deepSAD.pretrain(
|
||||
dataset,
|
||||
optimizer_name=cfg.settings["ae_optimizer_name"],
|
||||
lr=cfg.settings["ae_lr"],
|
||||
n_epochs=cfg.settings["ae_n_epochs"],
|
||||
lr_milestones=cfg.settings["ae_lr_milestone"],
|
||||
batch_size=cfg.settings["ae_batch_size"],
|
||||
weight_decay=cfg.settings["ae_weight_decay"],
|
||||
device=device,
|
||||
n_jobs_dataloader=n_jobs_dataloader,
|
||||
)
|
||||
|
||||
# Save pretraining results
|
||||
deepSAD.save_ae_results(export_json=xp_path + "/ae_results.json")
|
||||
|
||||
# Log training details
|
||||
logger.info("Training optimizer: %s" % cfg.settings["optimizer_name"])
|
||||
logger.info("Training learning rate: %g" % cfg.settings["lr"])
|
||||
logger.info("Training epochs: %d" % cfg.settings["n_epochs"])
|
||||
logger.info(
|
||||
"Training learning rate scheduler milestones: %s"
|
||||
% (cfg.settings["lr_milestone"],)
|
||||
)
|
||||
logger.info("Training batch size: %d" % cfg.settings["batch_size"])
|
||||
logger.info("Training weight decay: %g" % cfg.settings["weight_decay"])
|
||||
|
||||
# Train model on dataset
|
||||
deepSAD.train(
|
||||
dataset,
|
||||
optimizer_name=cfg.settings["ae_optimizer_name"],
|
||||
lr=cfg.settings["ae_lr"],
|
||||
n_epochs=cfg.settings["ae_n_epochs"],
|
||||
lr_milestones=cfg.settings["ae_lr_milestone"],
|
||||
batch_size=cfg.settings["ae_batch_size"],
|
||||
weight_decay=cfg.settings["ae_weight_decay"],
|
||||
optimizer_name=cfg.settings["optimizer_name"],
|
||||
lr=cfg.settings["lr"],
|
||||
n_epochs=cfg.settings["n_epochs"],
|
||||
lr_milestones=cfg.settings["lr_milestone"],
|
||||
batch_size=cfg.settings["batch_size"],
|
||||
weight_decay=cfg.settings["weight_decay"],
|
||||
device=device,
|
||||
n_jobs_dataloader=n_jobs_dataloader,
|
||||
)
|
||||
|
||||
# Save pretraining results
|
||||
deepSAD.save_ae_results(export_json=xp_path + "/ae_results.json")
|
||||
# Test model
|
||||
deepSAD.test(dataset, device=device, n_jobs_dataloader=n_jobs_dataloader)
|
||||
|
||||
# Log training details
|
||||
logger.info("Training optimizer: %s" % cfg.settings["optimizer_name"])
|
||||
logger.info("Training learning rate: %g" % cfg.settings["lr"])
|
||||
logger.info("Training epochs: %d" % cfg.settings["n_epochs"])
|
||||
logger.info(
|
||||
"Training learning rate scheduler milestones: %s"
|
||||
% (cfg.settings["lr_milestone"],)
|
||||
)
|
||||
logger.info("Training batch size: %d" % cfg.settings["batch_size"])
|
||||
logger.info("Training weight decay: %g" % cfg.settings["weight_decay"])
|
||||
# Save results, model, and configuration
|
||||
deepSAD.save_results(export_json=xp_path + "/results.json")
|
||||
deepSAD.save_model(export_model=xp_path + "/model.tar")
|
||||
cfg.save_config(export_json=xp_path + "/config.json")
|
||||
|
||||
# Train model on dataset
|
||||
deepSAD.train(
|
||||
dataset,
|
||||
optimizer_name=cfg.settings["optimizer_name"],
|
||||
lr=cfg.settings["lr"],
|
||||
n_epochs=cfg.settings["n_epochs"],
|
||||
lr_milestones=cfg.settings["lr_milestone"],
|
||||
batch_size=cfg.settings["batch_size"],
|
||||
weight_decay=cfg.settings["weight_decay"],
|
||||
device=device,
|
||||
n_jobs_dataloader=n_jobs_dataloader,
|
||||
)
|
||||
# Plot most anomalous and most normal test samples
|
||||
indices, labels, scores = zip(*deepSAD.results["test_scores"])
|
||||
indices, labels, scores = np.array(indices), np.array(labels), np.array(scores)
|
||||
idx_all_sorted = indices[np.argsort(scores)] # from lowest to highest score
|
||||
idx_normal_sorted = indices[labels == 0][
|
||||
np.argsort(scores[labels == 0])
|
||||
] # from lowest to highest score
|
||||
|
||||
# Test model
|
||||
deepSAD.test(dataset, device=device, n_jobs_dataloader=n_jobs_dataloader)
|
||||
if dataset_name in ("mnist", "fmnist", "cifar10", "elpv"):
|
||||
|
||||
# Save results, model, and configuration
|
||||
deepSAD.save_results(export_json=xp_path + "/results.json")
|
||||
deepSAD.save_model(export_model=xp_path + "/model.tar")
|
||||
cfg.save_config(export_json=xp_path + "/config.json")
|
||||
|
||||
# Plot most anomalous and most normal test samples
|
||||
indices, labels, scores = zip(*deepSAD.results["test_scores"])
|
||||
indices, labels, scores = np.array(indices), np.array(labels), np.array(scores)
|
||||
idx_all_sorted = indices[np.argsort(scores)] # from lowest to highest score
|
||||
idx_normal_sorted = indices[labels == 0][
|
||||
np.argsort(scores[labels == 0])
|
||||
] # from lowest to highest score
|
||||
|
||||
if dataset_name in ("mnist", "fmnist", "cifar10", "elpv"):
|
||||
|
||||
if dataset_name in ("mnist", "fmnist", "elpv"):
|
||||
X_all_low = dataset.test_set.data[idx_all_sorted[:32], ...].unsqueeze(1)
|
||||
X_all_high = dataset.test_set.data[idx_all_sorted[-32:], ...].unsqueeze(1)
|
||||
X_normal_low = dataset.test_set.data[idx_normal_sorted[:32], ...].unsqueeze(
|
||||
1
|
||||
)
|
||||
X_normal_high = dataset.test_set.data[
|
||||
idx_normal_sorted[-32:], ...
|
||||
].unsqueeze(1)
|
||||
|
||||
if dataset_name == "cifar10":
|
||||
X_all_low = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_all_sorted[:32], ...], (0, 3, 1, 2)
|
||||
if dataset_name in ("mnist", "fmnist", "elpv"):
|
||||
X_all_low = dataset.test_set.data[idx_all_sorted[:32], ...].unsqueeze(1)
|
||||
X_all_high = dataset.test_set.data[idx_all_sorted[-32:], ...].unsqueeze(
|
||||
1
|
||||
)
|
||||
)
|
||||
X_all_high = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_all_sorted[-32:], ...], (0, 3, 1, 2)
|
||||
)
|
||||
)
|
||||
X_normal_low = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_normal_sorted[:32], ...], (0, 3, 1, 2)
|
||||
)
|
||||
)
|
||||
X_normal_high = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_normal_sorted[-32:], ...], (0, 3, 1, 2)
|
||||
)
|
||||
)
|
||||
X_normal_low = dataset.test_set.data[
|
||||
idx_normal_sorted[:32], ...
|
||||
].unsqueeze(1)
|
||||
X_normal_high = dataset.test_set.data[
|
||||
idx_normal_sorted[-32:], ...
|
||||
].unsqueeze(1)
|
||||
|
||||
plot_images_grid(X_all_low, export_img=xp_path + "/all_low", padding=2)
|
||||
plot_images_grid(X_all_high, export_img=xp_path + "/all_high", padding=2)
|
||||
plot_images_grid(X_normal_low, export_img=xp_path + "/normals_low", padding=2)
|
||||
plot_images_grid(X_normal_high, export_img=xp_path + "/normals_high", padding=2)
|
||||
if dataset_name == "cifar10":
|
||||
X_all_low = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_all_sorted[:32], ...], (0, 3, 1, 2)
|
||||
)
|
||||
)
|
||||
X_all_high = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_all_sorted[-32:], ...], (0, 3, 1, 2)
|
||||
)
|
||||
)
|
||||
X_normal_low = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_normal_sorted[:32], ...], (0, 3, 1, 2)
|
||||
)
|
||||
)
|
||||
X_normal_high = torch.tensor(
|
||||
np.transpose(
|
||||
dataset.test_set.data[idx_normal_sorted[-32:], ...],
|
||||
(0, 3, 1, 2),
|
||||
)
|
||||
)
|
||||
|
||||
plot_images_grid(X_all_low, export_img=xp_path + "/all_low", padding=2)
|
||||
plot_images_grid(X_all_high, export_img=xp_path + "/all_high", padding=2)
|
||||
plot_images_grid(
|
||||
X_normal_low, export_img=xp_path + "/normals_low", padding=2
|
||||
)
|
||||
plot_images_grid(
|
||||
X_normal_high, export_img=xp_path + "/normals_high", padding=2
|
||||
)
|
||||
elif action == "infer":
|
||||
dataset = load_dataset(
|
||||
dataset_name,
|
||||
data_path,
|
||||
normal_class,
|
||||
known_outlier_class,
|
||||
n_known_outlier_classes,
|
||||
ratio_known_normal,
|
||||
ratio_known_outlier,
|
||||
ratio_pollution,
|
||||
random_state=np.random.RandomState(cfg.settings["seed"]),
|
||||
inference=True,
|
||||
)
|
||||
# Log random sample of known anomaly classes if more than 1 class
|
||||
if n_known_outlier_classes > 1:
|
||||
logger.info("Known anomaly classes: %s" % (dataset.known_outlier_classes,))
|
||||
|
||||
# Initialize DeepSAD model and set neural network phi
|
||||
deepSAD = DeepSAD(cfg.settings["eta"])
|
||||
deepSAD.set_network(net_name)
|
||||
|
||||
# If specified, load Deep SAD model (center c, network weights, and possibly autoencoder weights)
|
||||
if not load_model:
|
||||
logger.error(
|
||||
"For inference mode a model has to be loaded! Pass the --load_model option with the model path!"
|
||||
)
|
||||
return
|
||||
|
||||
deepSAD.load_model(model_path=load_model, load_ae=True, map_location=device)
|
||||
logger.info("Loading model from %s." % load_model)
|
||||
|
||||
inference_results = deepSAD.inference(
|
||||
dataset, device=device, n_jobs_dataloader=n_jobs_dataloader
|
||||
)
|
||||
inference_results_path = (
|
||||
Path(xp_path) / "inference" / Path(dataset.root).with_suffix(".npy").stem
|
||||
)
|
||||
inference_results_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
np.save(inference_results_path, inference_results, fix_imports=False)
|
||||
|
||||
logger.info(
|
||||
f"Inference: median={np.median(inference_results)} mean={np.mean(inference_results)} min={inference_results.min()} max={inference_results.max()}"
|
||||
)
|
||||
else:
|
||||
logger.error(f"Unknown action: {action}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user