133 lines
5.5 KiB
Python
Executable File
133 lines
5.5 KiB
Python
Executable File
import os
|
|
import sys
|
|
from datetime import datetime
|
|
|
|
import torch
|
|
import pickle
|
|
from tqdm import tqdm
|
|
|
|
path = os.path.abspath(__file__)
|
|
for i in range(2):
|
|
path = os.path.dirname(path)
|
|
PROJECT_ROOT = path
|
|
sys.path.append(PROJECT_ROOT)
|
|
|
|
from configs.config import ConfigManager
|
|
from datasets.dataset_factory import DatasetFactory
|
|
from modules.pipeline import Pipeline
|
|
from runners.runner import Runner
|
|
|
|
|
|
class InferenceEngine(Runner):
|
|
RESULTS_DIR_NAME: str = 'results'
|
|
LOG_DIR_NAME: str = 'log'
|
|
|
|
def __init__(self, config_path):
|
|
super().__init__(config_path)
|
|
|
|
''' Pipeline '''
|
|
self.pipeline_config = ConfigManager.get("settings", "pipeline")
|
|
self.pipeline = Pipeline(self.pipeline_config).to(self.device)
|
|
|
|
''' Experiment '''
|
|
self.model_path = ConfigManager.get("settings", "experiment", "model_path")
|
|
self.load_checkpoint(self.model_path)
|
|
self.load_experiment("inference")
|
|
|
|
''' Inference Results '''
|
|
self.inference_results_config = ConfigManager.get("settings", "results")
|
|
self.save_data_keys = self.inference_results_config["save_data_keys"]
|
|
self.save_output_keys = self.inference_results_config["save_output_keys"]
|
|
|
|
''' Test '''
|
|
self.test_config = ConfigManager.get("settings", "test")
|
|
self.test_dataset_config_list = self.test_config["dataset_list"]
|
|
self.test_set_list = []
|
|
seen_name = set()
|
|
for test_dataset_config in self.test_dataset_config_list:
|
|
if test_dataset_config["name"] not in seen_name:
|
|
seen_name.add(test_dataset_config["name"])
|
|
else:
|
|
raise ValueError("Duplicate test dataset name: {}".format(test_dataset_config["name"]))
|
|
test_set = DatasetFactory.create(test_dataset_config)
|
|
self.test_set_list.append(test_set)
|
|
del seen_name
|
|
|
|
self.print_info()
|
|
|
|
def run(self):
|
|
print("Inference start...")
|
|
self.test()
|
|
print("Inference finished!")
|
|
|
|
def test(self):
|
|
self.pipeline.eval()
|
|
with torch.no_grad():
|
|
for dataset_idx, test_set in enumerate(self.test_set_list):
|
|
test_set_name = self.test_dataset_config_list[dataset_idx]["name"]
|
|
ratio = self.test_dataset_config_list[dataset_idx]["ratio"]
|
|
|
|
test_loader = test_set.get_loader()
|
|
loop = tqdm(enumerate(test_loader), total=int(len(test_loader)))
|
|
for i, data in loop:
|
|
test_set.process_batch(data, self.device)
|
|
output = self.pipeline(data, Pipeline.TEST_MODE)
|
|
self.save_output(output, data, test_set_name, i)
|
|
loop.set_description(
|
|
f'Inference (Test: {test_set_name}, ratio={ratio})')
|
|
|
|
def save_output(self, output, data, test_set_name, idx):
|
|
results_dir = os.path.join(str(self.experiment_path), InferenceEngine.RESULTS_DIR_NAME)
|
|
if not os.path.exists(os.path.join(results_dir,test_set_name)):
|
|
os.makedirs(os.path.join(results_dir,test_set_name))
|
|
save_path = os.path.join(results_dir, test_set_name, f"{idx}.pkl")
|
|
data = {key: value for key, value in data.items() if key in self.save_data_keys}
|
|
output = {key: value for key, value in output.items() if key in self.save_output_keys}
|
|
output_converted = {key: value.cpu().numpy() if torch.is_tensor(value) else value for key, value in output.items()}
|
|
data_converted = {key: value.cpu().numpy() if torch.is_tensor(value) else value for key, value in data.items()}
|
|
with open(save_path, "wb") as f:
|
|
pickle.dump({"output":output_converted,"data":data_converted}, f)
|
|
|
|
def load_checkpoint(self, model_path):
|
|
self.pipeline.load(model_path)
|
|
print(f"Checkpoint loaded from {model_path}")
|
|
|
|
def load_experiment(self, backup_name=None):
|
|
super().load_experiment(backup_name)
|
|
|
|
def create_experiment(self, backup_name=None):
|
|
super().create_experiment(backup_name)
|
|
results_dir = os.path.join(str(self.experiment_path), InferenceEngine.RESULTS_DIR_NAME)
|
|
os.makedirs(results_dir)
|
|
|
|
|
|
def print_info(self):
|
|
def print_dataset(config, dataset):
|
|
print("\t name: {}".format(config["name"]))
|
|
print("\t source: {}".format(config["source"]))
|
|
print("\t data_type: {}".format(config["data_type"]))
|
|
print("\t total_length: {}".format(len(dataset)))
|
|
print("\t ratio: {}".format(config["ratio"]))
|
|
print()
|
|
|
|
super().print_info()
|
|
table_size = 70
|
|
print(f"{'+' + '-' * (table_size // 2)} Pipeline {'-' * (table_size // 2)}" + '+')
|
|
print(self.pipeline)
|
|
print(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+')
|
|
for i, test_dataset_config in enumerate(self.test_dataset_config_list):
|
|
print(f"test dataset {i}: ")
|
|
print_dataset(test_dataset_config, self.test_set_list[i])
|
|
|
|
print(f"{'+' + '-' * (table_size // 2)}----------{'-' * (table_size // 2)}" + '+')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=str, default="configs/local_inference_config.yaml")
|
|
args = parser.parse_args()
|
|
infenrence_engine = InferenceEngine(args.config)
|
|
infenrence_engine.run()
|