diff --git a/abl/bridge/base_bridge.py b/abl/bridge/base_bridge.py index 03054f7..869ea39 100644 --- a/abl/bridge/base_bridge.py +++ b/abl/bridge/base_bridge.py @@ -1,52 +1,64 @@ from abc import ABCMeta, abstractmethod -from typing import Any, List, Tuple +from typing import Any, List, Optional, Tuple, Union from ..learning import ABLModel from ..reasoning import ReasonerBase +from ..structures import ListData +DataSet = Tuple[List[List[Any]], Optional[List[List[Any]]], List[List[Any]]] -class BaseBridge(metaclass=ABCMeta): +class BaseBridge(metaclass=ABCMeta): def __init__(self, model: ABLModel, abducer: ReasonerBase) -> None: if not isinstance(model, ABLModel): - raise TypeError("Expected an ABLModel") + raise TypeError( + "Expected an instance of ABLModel, but received type: {}".format( + type(model) + ) + ) if not isinstance(abducer, ReasonerBase): - raise TypeError("Expected an ReasonerBase") - + raise TypeError( + "Expected an instance of ReasonerBase, but received type: {}".format( + type(abducer) + ) + ) + self.model = model self.abducer = abducer @abstractmethod - def predict(self, X: List[List[Any]]) -> Tuple[List[List[Any]], List[List[Any]]]: + def predict( + self, data_samples: ListData + ) -> Tuple[List[List[Any]], List[List[Any]]]: """Placeholder for predict labels from input.""" pass @abstractmethod - def abduce_pseudo_label(self, pseudo_label: List[List[Any]]) -> List[List[Any]]: + def abduce_pseudo_label(self, data_samples: ListData) -> List[List[Any]]: """Placeholder for abduce pseudo labels.""" + pass @abstractmethod - def idx_to_pseudo_label(self, idx: List[List[Any]]) -> List[List[Any]]: + def idx_to_pseudo_label(self, data_samples: ListData) -> List[List[Any]]: """Placeholder for map label space to symbol space.""" pass @abstractmethod - def pseudo_label_to_idx(self, pseudo_label: List[List[Any]]) -> List[List[Any]]: + def pseudo_label_to_idx(self, data_samples: ListData) -> List[List[Any]]: """Placeholder for map symbol space to label space.""" pass - + @abstractmethod - def train(self, train_data): + def train(self, train_data: Union[ListData, DataSet]): """Placeholder for train loop of ABductive Learning.""" pass @abstractmethod - def test(self, test_data): + def valid(self, valid_data: Union[ListData, DataSet]) -> None: """Placeholder for model test.""" pass @abstractmethod - def valid(self, valid_data): + def test(self, test_data: Union[ListData, DataSet]) -> None: """Placeholder for model validation.""" pass - \ No newline at end of file diff --git a/abl/bridge/simple_bridge.py b/abl/bridge/simple_bridge.py index 3aecffd..9093bc1 100644 --- a/abl/bridge/simple_bridge.py +++ b/abl/bridge/simple_bridge.py @@ -1,13 +1,14 @@ -from ..learning import ABLModel -from ..reasoning import ReasonerBase -from ..evaluation import BaseMetric -from .base_bridge import BaseBridge -from typing import List, Union, Any, Tuple, Dict, Optional +import os.path as osp +from typing import Any, Dict, List, Optional, Tuple, Union + from numpy import ndarray -from torch.utils.data import DataLoader -from ..dataset import BridgeDataset -from ..utils.logger import print_log +from ..evaluation import BaseMetric +from ..learning import ABLModel +from ..reasoning import ReasonerBase +from ..structures import ListData +from ..utils import print_log +from .base_bridge import BaseBridge, DataSet class SimpleBridge(BaseBridge): @@ -20,85 +21,99 @@ class SimpleBridge(BaseBridge): super().__init__(model, abducer) self.metric_list = metric_list - def predict(self, X) -> Tuple[List[List[Any]], ndarray]: - pred_res = self.model.predict(X) - pred_idx, pred_prob = pred_res["label"], pred_res["prob"] - return pred_idx, pred_prob - + # TODO: add abducer.mapping to the property of SimpleBridge + + def predict(self, data_samples: ListData) -> Tuple[List[ndarray], List[ndarray]]: + self.model.predict(data_samples) + return data_samples["pred_idx"], data_samples.get("pred_prob", None) + def abduce_pseudo_label( self, - pred_prob: ndarray, - pred_pseudo_label: List[List[Any]], - Y: List[Any], + data_samples: ListData, max_revision: int = -1, require_more_revision: int = 0, ) -> List[List[Any]]: - return self.abducer.batch_abduce(pred_prob, pred_pseudo_label, Y, max_revision, require_more_revision) + self.abducer.batch_abduce(data_samples, max_revision, require_more_revision) + return data_samples["abduced_pseudo_label"] def idx_to_pseudo_label( - self, idx: List[List[Any]], mapping: Dict = None + self, data_samples: ListData, mapping: Optional[Dict] = None ) -> List[List[Any]]: if mapping is None: mapping = self.abducer.mapping - return [[mapping[_idx] for _idx in sub_list] for sub_list in idx] + pred_idx = data_samples.pred_idx + data_samples.pred_pseudo_label = [ + [mapping[_idx] for _idx in sub_list] for sub_list in pred_idx + ] + return data_samples["pred_pseudo_label"] def pseudo_label_to_idx( - self, pseudo_label: List[List[Any]], mapping: Dict = None + self, data_samples: ListData, mapping: Optional[Dict] = None ) -> List[List[Any]]: if mapping is None: mapping = self.abducer.remapping - return [ - [mapping[_pseudo_label] for _pseudo_label in sub_list] - for sub_list in pseudo_label + abduced_idx = [ + [mapping[_abduced_pseudo_label] for _abduced_pseudo_label in sub_list] + for sub_list in data_samples.abduced_pseudo_label ] + data_samples.abduced_idx = abduced_idx + return data_samples["abduced_idx"] + + def data_preprocess(self, X: List[Any], gt_pseudo_label: List[Any], Y: List[Any]) -> ListData: + data_samples = ListData() + + data_samples.X = X + data_samples.gt_pseudo_label = gt_pseudo_label + data_samples.Y = Y + + return data_samples def train( self, - train_data: Tuple[List[List[Any]], Optional[List[List[Any]]], List[List[Any]]], - epochs: int = 50, - batch_size: Union[int, float] = -1, + train_data: Union[ListData, DataSet], + loops: int = 50, + segment_size: Union[int, float] = -1, eval_interval: int = 1, + save_interval: Optional[int] = None, + save_dir: Optional[str] = None, ): - dataset = BridgeDataset(*train_data) - data_loader = DataLoader( - dataset, - batch_size=batch_size, - collate_fn=lambda data_list: [list(data) for data in zip(*data_list)], - ) - - for epoch in range(epochs): - for seg_idx, (X, Z, Y) in enumerate(data_loader): - pred_idx, pred_prob = self.predict(X) - pred_pseudo_label = self.idx_to_pseudo_label(pred_idx) - abduced_pseudo_label = self.abduce_pseudo_label( - pred_prob, pred_pseudo_label, Y - ) - abduced_label = self.pseudo_label_to_idx(abduced_pseudo_label) - loss = self.model.train(X, abduced_label) + if isinstance(train_data, ListData): + data_samples = train_data + else: + data_samples = self.data_preprocess(*train_data) + + for loop in range(loops): + for seg_idx in range((len(data_samples) - 1) // segment_size + 1): + sub_data_samples = data_samples[ + seg_idx * segment_size : (seg_idx + 1) * segment_size + ] + self.predict(sub_data_samples) + self.idx_to_pseudo_label(sub_data_samples) + self.abduce_pseudo_label(sub_data_samples) + self.pseudo_label_to_idx(sub_data_samples) + loss = self.model.train(sub_data_samples) print_log( - f"Epoch(train) [{epoch + 1}] [{(seg_idx + 1):3}/{len(data_loader)}] model loss is {loss:.5f}", + f"loop(train) [{loop + 1}/{loops}] segment(train) [{(seg_idx + 1)}/{(len(data_samples) - 1) // segment_size + 1}] model loss is {loss:.5f}", logger="current", ) - if (epoch + 1) % eval_interval == 0 or epoch == epochs - 1: - print_log(f"Evaluation start: Epoch(val) [{epoch}]", logger="current") + if (loop + 1) % eval_interval == 0 or loop == loops - 1: + print_log(f"Evaluation start: loop(val) [{loop + 1}]", logger="current") self.valid(train_data) - def _valid(self, data_loader): - for X, Z, Y in data_loader: - pred_idx, pred_prob = self.predict(X) - pred_pseudo_label = self.idx_to_pseudo_label(pred_idx) - data_samples = dict( - pred_idx=pred_idx, - pred_prob=pred_prob, - pred_pseudo_label=pred_pseudo_label, - gt_pseudo_label=Z, - Y=Y, - logic_forward=self.abducer.kb.logic_forward, - ) + if save_interval is not None and ((loop + 1) % save_interval == 0 or loop == loops - 1): + print_log(f"Saving model: loop(save) [{loop + 1}]", logger="current") + self.model.save(save_path=osp.join(save_dir, f"model_checkpoint_loop_{loop + 1}.pth")) + + def _valid(self, data_samples: ListData, batch_size: int = 128) -> None: + for seg_idx in range((len(data_samples) - 1) // batch_size + 1): + sub_data_samples = data_samples[seg_idx * batch_size : (seg_idx + 1) * batch_size] + self.predict(sub_data_samples) + self.idx_to_pseudo_label(sub_data_samples) + for metric in self.metric_list: - metric.process(data_samples) + metric.process(sub_data_samples) res = dict() for metric in self.metric_list: @@ -108,14 +123,12 @@ class SimpleBridge(BaseBridge): msg += k + f": {v:.3f} " print_log(msg, logger="current") - def valid(self, valid_data, batch_size=1000): - dataset = BridgeDataset(*valid_data) - data_loader = DataLoader( - dataset, - batch_size=batch_size, - collate_fn=lambda data_list: [list(data) for data in zip(*data_list)], - ) - self._valid(data_loader) - - def test(self, test_data, batch_size=1000): - self.valid(test_data, batch_size) + def valid(self, valid_data: Union[ListData, DataSet], batch_size: int = 128) -> None: + if not isinstance(valid_data, ListData): + data_samples = self.data_preprocess(*valid_data) + else: + data_samples = valid_data + self._valid(data_samples, batch_size=batch_size) + + def test(self, test_data: Union[ListData, DataSet], batch_size: int = 128) -> None: + self.valid(test_data, batch_size=batch_size) diff --git a/abl/dataset/__init__.py b/abl/dataset/__init__.py index 6be0df1..a487476 100644 --- a/abl/dataset/__init__.py +++ b/abl/dataset/__init__.py @@ -1,3 +1,4 @@ from .bridge_dataset import BridgeDataset from .classification_dataset import ClassificationDataset -from .regression_dataset import RegressionDataset \ No newline at end of file +from .prediction_dataset import PredictionDataset +from .regression_dataset import RegressionDataset diff --git a/abl/dataset/bridge_dataset.py b/abl/dataset/bridge_dataset.py index bb0ce98..a7d32c5 100644 --- a/abl/dataset/bridge_dataset.py +++ b/abl/dataset/bridge_dataset.py @@ -1,5 +1,6 @@ +from typing import Any, List, Tuple + from torch.utils.data import Dataset -from typing import List, Any, Tuple class BridgeDataset(Dataset): diff --git a/abl/dataset/classification_dataset.py b/abl/dataset/classification_dataset.py index 28f9299..1663642 100644 --- a/abl/dataset/classification_dataset.py +++ b/abl/dataset/classification_dataset.py @@ -1,6 +1,7 @@ +from typing import Any, Callable, List, Tuple + import torch from torch.utils.data import Dataset -from typing import List, Any, Tuple, Callable class ClassificationDataset(Dataset): diff --git a/abl/dataset/prediction_dataset.py b/abl/dataset/prediction_dataset.py new file mode 100644 index 0000000..8e3c717 --- /dev/null +++ b/abl/dataset/prediction_dataset.py @@ -0,0 +1,56 @@ +from typing import Any, Callable, List, Tuple + +import torch +from torch.utils.data import Dataset + + +class PredictionDataset(Dataset): + def __init__(self, X: List[Any], transform: Callable[..., Any] = None): + """ + Initialize the dataset used for classification task. + + Parameters + ---------- + X : List[Any] + The input data. + transform : Callable[..., Any], optional + A function/transform that takes in an object and returns a transformed version. Defaults to None. + """ + if not isinstance(X, list): + raise ValueError("X should be of type list.") + + self.X = X + self.transform = transform + + def __len__(self) -> int: + """ + Return the length of the dataset. + + Returns + ------- + int + The length of the dataset. + """ + return len(self.X) + + def __getitem__(self, index: int) -> Tuple[Any, torch.Tensor]: + """ + Get the item at the given index. + + Parameters + ---------- + index : int + The index of the item to get. + + Returns + ------- + Tuple[Any, torch.Tensor] + A tuple containing the object and its label. + """ + if index >= len(self): + raise ValueError("index range error") + + x = self.X[index] + if self.transform is not None: + x = self.transform(x) + return x diff --git a/abl/dataset/regression_dataset.py b/abl/dataset/regression_dataset.py index 8cf136c..118ac65 100644 --- a/abl/dataset/regression_dataset.py +++ b/abl/dataset/regression_dataset.py @@ -1,6 +1,7 @@ +from typing import Any, List, Tuple + import torch from torch.utils.data import Dataset -from typing import List, Any, Tuple class RegressionDataset(Dataset): diff --git a/abl/evaluation/__init__.py b/abl/evaluation/__init__.py index a849d68..3106412 100644 --- a/abl/evaluation/__init__.py +++ b/abl/evaluation/__init__.py @@ -1,3 +1,3 @@ from .base_metric import BaseMetric -from .symbol_metric import SymbolMetric from .semantics_metric import SemanticsMetric +from .symbol_metric import SymbolMetric diff --git a/abl/evaluation/base_metric.py b/abl/evaluation/base_metric.py index 44364f8..e18f452 100644 --- a/abl/evaluation/base_metric.py +++ b/abl/evaluation/base_metric.py @@ -1,8 +1,8 @@ +import logging from abc import ABCMeta, abstractmethod from typing import Any, List, Optional, Sequence -from ..utils import print_log -import logging +from ..utils import print_log class BaseMetric(metaclass=ABCMeta): diff --git a/abl/evaluation/semantics_metric.py b/abl/evaluation/semantics_metric.py index 3333daf..718cfea 100644 --- a/abl/evaluation/semantics_metric.py +++ b/abl/evaluation/semantics_metric.py @@ -1,25 +1,22 @@ from typing import Optional, Sequence + +from ..reasoning import BaseKB from .base_metric import BaseMetric -class ABLMetric(): - pass class SemanticsMetric(BaseMetric): - def __init__(self, prefix: Optional[str] = None) -> None: + def __init__(self, kb: BaseKB = None, prefix: Optional[str] = None) -> None: super().__init__(prefix) + self.kb = kb def process(self, data_samples: Sequence[dict]) -> None: - pred_pseudo_label = data_samples["pred_pseudo_label"] - gt_Y = data_samples["Y"] - logic_forward = data_samples["logic_forward"] - - for pred_z, y in zip(pred_pseudo_label, gt_Y): - if logic_forward(pred_z) == y: + for data_sample in data_samples: + if self.kb.check_equal(data_sample, data_sample["Y"][0]): self.results.append(1) else: self.results.append(0) - + def compute_metrics(self, results: list) -> dict: metrics = dict() metrics["semantics_accuracy"] = sum(results) / len(results) - return metrics \ No newline at end of file + return metrics diff --git a/abl/evaluation/symbol_metric.py b/abl/evaluation/symbol_metric.py index 3c0c216..c2d7938 100644 --- a/abl/evaluation/symbol_metric.py +++ b/abl/evaluation/symbol_metric.py @@ -1,4 +1,5 @@ -from typing import Optional, Sequence, Callable +from typing import Optional, Sequence + from .base_metric import BaseMetric diff --git a/abl/learning/abl_model.py b/abl/learning/abl_model.py index f1512fe..6685cc4 100644 --- a/abl/learning/abl_model.py +++ b/abl/learning/abl_model.py @@ -10,8 +10,10 @@ # # ================================================================# import pickle -from utils import flatten, reform_idx -from typing import List, Any, Optional +from typing import Any, Dict + +from ..structures import ListData +from ..utils import reform_idx class ABLModel: @@ -30,7 +32,7 @@ class ABLModel: Methods ------- - predict(X: List[List[Any]], mapping: Optional[dict] = None) -> dict + predict(X: List[List[Any]], mapping: Optional[Dict] = None) -> Dict Predict the labels and probabilities for the given data. valid(X: List[List[Any]], Y: List[Any]) -> float Calculate the accuracy score for the given data. @@ -42,20 +44,13 @@ class ABLModel: Load the model from a file. """ - def __init__(self, base_model) -> None: - self.classifier_list = [] - self.classifier_list.append(base_model) + def __init__(self, base_model: Any) -> None: + if not (hasattr(base_model, "fit") and hasattr(base_model, "predict")): + raise NotImplementedError("The base_model should implement fit and predict methods.") - if not ( - hasattr(base_model, "fit") - and hasattr(base_model, "predict") - and hasattr(base_model, "score") - ): - raise NotImplementedError( - "base_model should have fit, predict and score methods." - ) + self.base_model = base_model - def predict(self, X: List[List[Any]], mapping: Optional[dict] = None) -> dict: + def predict(self, data_samples: ListData) -> Dict: """ Predict the labels and probabilities for the given data. @@ -63,53 +58,30 @@ class ABLModel: ---------- X : List[List[Any]] The data to predict on. - mapping : Optional[dict], optional - A mapping dictionary to map labels to their original values, by default None. Returns ------- dict A dictionary containing the predicted labels and probabilities. """ - model = self.classifier_list[0] - data_X = flatten(X) + model = self.base_model + data_X = data_samples.flatten("X") if hasattr(model, "predict_proba"): prob = model.predict_proba(X=data_X) label = prob.argmax(axis=1) - prob = reform_idx(prob, X) + prob = reform_idx(prob, data_samples["X"]) else: prob = None label = model.predict(X=data_X) + label = reform_idx(label, data_samples["X"]) - if mapping is not None: - label = [mapping[y] for y in label] - - label = reform_idx(label, X) + data_samples.pred_idx = label + if prob is not None: + data_samples.pred_prob = prob return {"label": label, "prob": prob} - def valid(self, X: List[List[Any]], Y: List[Any]) -> float: - """ - Calculate the accuracy for the given data. - - Parameters - ---------- - X : List[List[Any]] - The data to calculate the accuracy on. - Y : List[Any] - The true labels for the given data. - - Returns - ------- - float - The accuracy score for the given data. - """ - data_X = flatten(X) - data_Y = flatten(Y) - score = self.classifier_list[0].score(X=data_X, y=data_Y) - return score - - def train(self, X: List[List[Any]], Y: List[Any]) -> float: + def train(self, data_samples: ListData) -> float: """ Train the model on the given data. @@ -125,29 +97,30 @@ class ABLModel: float The loss value of the trained model. """ - data_X = flatten(X) - data_Y = flatten(Y) - return self.classifier_list[0].fit(X=data_X, y=data_Y) + data_X = data_samples.flatten("X") + data_y = data_samples.flatten("abduced_idx") + return self.base_model.fit(X=data_X, y=data_y) def _model_operation(self, operation: str, *args, **kwargs): - model = self.classifier_list[0] + model = self.base_model if hasattr(model, operation): method = getattr(model, operation) method(*args, **kwargs) else: - try: - if not f"{operation}_path" in kwargs.keys(): - raise ValueError(f"'{operation}_path' should not be None") - if operation == "save": - with open(kwargs["save_path"], 'wb') as file: - pickle.dump(model, file, protocol=pickle.HIGHEST_PROTOCOL) - elif operation == "load": - with open(kwargs["load_path"], 'rb') as file: - self.classifier_list[0] = pickle.load(file) - except: - raise NotImplementedError( - f"{type(model).__name__} object doesn't have the {operation} method" - ) + if not f"{operation}_path" in kwargs.keys(): + raise ValueError(f"'{operation}_path' should not be None") + else: + try: + if operation == "save": + with open(kwargs["save_path"], "wb") as file: + pickle.dump(model, file, protocol=pickle.HIGHEST_PROTOCOL) + elif operation == "load": + with open(kwargs["load_path"], "rb") as file: + self.base_model = pickle.load(file) + except: + raise NotImplementedError( + f"{type(model).__name__} object doesn't have the {operation} method and the default pickle-based {operation} method failed." + ) def save(self, *args, **kwargs) -> None: """ diff --git a/abl/learning/basic_nn.py b/abl/learning/basic_nn.py index b9b0b36..b1da93c 100644 --- a/abl/learning/basic_nn.py +++ b/abl/learning/basic_nn.py @@ -10,14 +10,16 @@ # # ================================================================# -import torch +import os +import logging +from typing import Any, Callable, List, Optional, T, Tuple + import numpy +import torch from torch.utils.data import DataLoader -from ..utils.logger import print_log -from ..dataset import ClassificationDataset -import os -from typing import List, Any, T, Optional, Callable, Tuple +from ..dataset import ClassificationDataset, PredictionDataset +from ..utils.logger import print_log class BasicNN: @@ -99,9 +101,7 @@ class BasicNN: loss_value = self.train_epoch(data_loader) if self.save_interval is not None and (epoch + 1) % self.save_interval == 0: if self.save_dir is None: - raise ValueError( - "save_dir should not be None if save_interval is not None." - ) + raise ValueError("save_dir should not be None if save_interval is not None.") self.save(epoch + 1) if self.stop_loss is not None and loss_value < self.stop_loss: break @@ -191,7 +191,7 @@ class BasicNN: with torch.no_grad(): results = [] - for data, _ in data_loader: + for data in data_loader: data = data.to(device) out = model(data) results.append(out) @@ -199,7 +199,10 @@ class BasicNN: return torch.cat(results, axis=0) def predict( - self, data_loader: DataLoader = None, X: List[Any] = None + self, + data_loader: DataLoader = None, + X: List[Any] = None, + test_transform: Callable[..., Any] = None, ) -> numpy.ndarray: """ Predict the class of the input data. @@ -218,11 +221,28 @@ class BasicNN: """ if data_loader is None: - data_loader = self._data_loader(X) + if test_transform is None: + print_log( + "Transform used in the training phase will be used in prediction.", + "current", + level=logging.WARNING, + ) + dataset = PredictionDataset(X, self.transform) + else: + dataset = PredictionDataset(X, test_transform) + data_loader = DataLoader( + dataset, + batch_size=self.batch_size, + num_workers=int(self.num_workers), + collate_fn=self.collate_fn, + ) return self._predict(data_loader).argmax(axis=1).cpu().numpy() def predict_proba( - self, data_loader: DataLoader = None, X: List[Any] = None + self, + data_loader: DataLoader = None, + X: List[Any] = None, + test_transform: Callable[..., Any] = None, ) -> numpy.ndarray: """ Predict the probability of each class for the input data. @@ -241,7 +261,21 @@ class BasicNN: """ if data_loader is None: - data_loader = self._data_loader(X) + if test_transform is None: + print_log( + "Transform used in the training phase will be used in prediction.", + "current", + level=logging.WARNING, + ) + dataset = PredictionDataset(X, self.transform) + else: + dataset = PredictionDataset(X, test_transform) + data_loader = DataLoader( + dataset, + batch_size=self.batch_size, + num_workers=int(self.num_workers), + collate_fn=self.collate_fn, + ) return self._predict(data_loader).softmax(axis=1).cpu().numpy() def _score(self, data_loader) -> Tuple[float, float]: @@ -313,15 +347,14 @@ class BasicNN: if data_loader is None: data_loader = self._data_loader(X, y) mean_loss, accuracy = self._score(data_loader) - print_log( - f"mean loss: {mean_loss:.3f}, accuray: {accuracy:.3f}", logger="current" - ) + print_log(f"mean loss: {mean_loss:.3f}, accuray: {accuracy:.3f}", logger="current") return accuracy def _data_loader( self, X: List[Any], y: List[int] = None, + shuffle: bool = True, ) -> DataLoader: """ Generate a DataLoader for user-provided input and target data. @@ -350,7 +383,7 @@ class BasicNN: data_loader = DataLoader( dataset, batch_size=self.batch_size, - shuffle=True, + shuffle=shuffle, num_workers=int(self.num_workers), collate_fn=self.collate_fn, ) @@ -368,14 +401,13 @@ class BasicNN: The path to save the model, by default None. """ if self.save_dir is None and save_path is None: - raise ValueError( - "'save_dir' and 'save_path' should not be None simultaneously." - ) + raise ValueError("'save_dir' and 'save_path' should not be None simultaneously.") - if save_path is None: - save_path = os.path.join( - self.save_dir, f"model_checkpoint_epoch_{epoch_id}.pth" - ) + if save_path is not None: + if not os.path.exists(os.path.dirname(save_path)): + os.makedirs(os.path.dirname(save_path)) + else: + save_path = os.path.join(self.save_dir, f"model_checkpoint_epoch_{epoch_id}.pth") if not os.path.exists(self.save_dir): os.makedirs(self.save_dir) diff --git a/abl/structures/__init__.py b/abl/structures/__init__.py new file mode 100644 index 0000000..52b5af3 --- /dev/null +++ b/abl/structures/__init__.py @@ -0,0 +1,2 @@ +from .base_data_element import BaseDataElement +from .list_data import ListData \ No newline at end of file diff --git a/abl/structures/base_data_element.py b/abl/structures/base_data_element.py new file mode 100644 index 0000000..03176d1 --- /dev/null +++ b/abl/structures/base_data_element.py @@ -0,0 +1,629 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import copy +from typing import Any, Iterator, Optional, Tuple, Type, Union + +import numpy as np +import torch + + +class BaseDataElement: + """A base data interface that supports Tensor-like and dict-like + operations. + + A typical data elements refer to predicted results or ground truth labels + on a task, such as predicted bboxes, instance masks, semantic + segmentation masks, etc. Because groundtruth labels and predicted results + often have similar properties (for example, the predicted bboxes and the + groundtruth bboxes), MMEngine uses the same abstract data interface to + encapsulate predicted results and groundtruth labels, and it is recommended + to use different name conventions to distinguish them, such as using + ``gt_instances`` and ``pred_instances`` to distinguish between labels and + predicted results. Additionally, we distinguish data elements at instance + level, pixel level, and label level. Each of these types has its own + characteristics. Therefore, MMEngine defines the base class + ``BaseDataElement``, and implement ``InstanceData``, ``PixelData``, and + ``LabelData`` inheriting from ``BaseDataElement`` to represent different + types of ground truth labels or predictions. + + Another common data element is sample data. A sample data consists of input + data (such as an image) and its annotations and predictions. In general, + an image can have multiple types of annotations and/or predictions at the + same time (for example, both pixel-level semantic segmentation annotations + and instance-level detection bboxes annotations). All labels and + predictions of a training sample are often passed between Dataset, Model, + Visualizer, and Evaluator components. In order to simplify the interface + between components, we can treat them as a large data element and + encapsulate them. Such data elements are generally called XXDataSample in + the OpenMMLab. Therefore, Similar to `nn.Module`, the `BaseDataElement` + allows `BaseDataElement` as its attribute. Such a class generally + encapsulates all the data of a sample in the algorithm library, and its + attributes generally are various types of data elements. For example, + MMDetection is assigned by the BaseDataElement to encapsulate all the data + elements of the sample labeling and prediction of a sample in the + algorithm library. + + The attributes in ``BaseDataElement`` are divided into two parts, + the ``metainfo`` and the ``data`` respectively. + + - ``metainfo``: Usually contains the + information about the image such as filename, + image_shape, pad_shape, etc. The attributes can be accessed or + modified by dict-like or object-like operations, such as + ``.`` (for data access and modification), ``in``, ``del``, + ``pop(str)``, ``get(str)``, ``metainfo_keys()``, + ``metainfo_values()``, ``metainfo_items()``, ``set_metainfo()`` (for + set or change key-value pairs in metainfo). + + - ``data``: Annotations or model predictions are + stored. The attributes can be accessed or modified by + dict-like or object-like operations, such as + ``.``, ``in``, ``del``, ``pop(str)``, ``get(str)``, ``keys()``, + ``values()``, ``items()``. Users can also apply tensor-like + methods to all :obj:`torch.Tensor` in the ``data_fields``, + such as ``.cuda()``, ``.cpu()``, ``.numpy()``, ``.to()``, + ``to_tensor()``, ``.detach()``. + + Args: + metainfo (dict, optional): A dict contains the meta information + of single image, such as ``dict(img_shape=(512, 512, 3), + scale_factor=(1, 1, 1, 1))``. Defaults to None. + kwargs (dict, optional): A dict contains annotations of single image or + model predictions. Defaults to None. + + Examples: + >>> import torch + >>> from mmengine.structures import BaseDataElement + >>> gt_instances = BaseDataElement() + >>> bboxes = torch.rand((5, 4)) + >>> scores = torch.rand((5,)) + >>> img_id = 0 + >>> img_shape = (800, 1333) + >>> gt_instances = BaseDataElement( + ... metainfo=dict(img_id=img_id, img_shape=img_shape), + ... bboxes=bboxes, scores=scores) + >>> gt_instances = BaseDataElement( + ... metainfo=dict(img_id=img_id, img_shape=(640, 640))) + + >>> # new + >>> gt_instances1 = gt_instances.new( + ... metainfo=dict(img_id=1, img_shape=(640, 640)), + ... bboxes=torch.rand((5, 4)), + ... scores=torch.rand((5,))) + >>> gt_instances2 = gt_instances1.new() + + >>> # add and process property + >>> gt_instances = BaseDataElement() + >>> gt_instances.set_metainfo(dict(img_id=9, img_shape=(100, 100))) + >>> assert 'img_shape' in gt_instances.metainfo_keys() + >>> assert 'img_shape' in gt_instances + >>> assert 'img_shape' not in gt_instances.keys() + >>> assert 'img_shape' in gt_instances.all_keys() + >>> print(gt_instances.img_shape) + (100, 100) + >>> gt_instances.scores = torch.rand((5,)) + >>> assert 'scores' in gt_instances.keys() + >>> assert 'scores' in gt_instances + >>> assert 'scores' in gt_instances.all_keys() + >>> assert 'scores' not in gt_instances.metainfo_keys() + >>> print(gt_instances.scores) + tensor([0.5230, 0.7885, 0.2426, 0.3911, 0.4876]) + >>> gt_instances.bboxes = torch.rand((5, 4)) + >>> assert 'bboxes' in gt_instances.keys() + >>> assert 'bboxes' in gt_instances + >>> assert 'bboxes' in gt_instances.all_keys() + >>> assert 'bboxes' not in gt_instances.metainfo_keys() + >>> print(gt_instances.bboxes) + tensor([[0.0900, 0.0424, 0.1755, 0.4469], + [0.8648, 0.0592, 0.3484, 0.0913], + [0.5808, 0.1909, 0.6165, 0.7088], + [0.5490, 0.4209, 0.9416, 0.2374], + [0.3652, 0.1218, 0.8805, 0.7523]]) + + >>> # delete and change property + >>> gt_instances = BaseDataElement( + ... metainfo=dict(img_id=0, img_shape=(640, 640)), + ... bboxes=torch.rand((6, 4)), scores=torch.rand((6,))) + >>> gt_instances.set_metainfo(dict(img_shape=(1280, 1280))) + >>> gt_instances.img_shape # (1280, 1280) + >>> gt_instances.bboxes = gt_instances.bboxes * 2 + >>> gt_instances.get('img_shape', None) # (1280, 1280) + >>> gt_instances.get('bboxes', None) # 6x4 tensor + >>> del gt_instances.img_shape + >>> del gt_instances.bboxes + >>> assert 'img_shape' not in gt_instances + >>> assert 'bboxes' not in gt_instances + >>> gt_instances.pop('img_shape', None) # None + >>> gt_instances.pop('bboxes', None) # None + + >>> # Tensor-like + >>> cuda_instances = gt_instances.cuda() + >>> cuda_instances = gt_instances.to('cuda:0') + >>> cpu_instances = cuda_instances.cpu() + >>> cpu_instances = cuda_instances.to('cpu') + >>> fp16_instances = cuda_instances.to( + ... device=None, dtype=torch.float16, non_blocking=False, + ... copy=False, memory_format=torch.preserve_format) + >>> cpu_instances = cuda_instances.detach() + >>> np_instances = cpu_instances.numpy() + + >>> # print + >>> metainfo = dict(img_shape=(800, 1196, 3)) + >>> gt_instances = BaseDataElement( + ... metainfo=metainfo, det_labels=torch.LongTensor([0, 1, 2, 3])) + >>> sample = BaseDataElement(metainfo=metainfo, + ... gt_instances=gt_instances) + >>> print(sample) + + ) at 0x7f0fea49e130> + + >>> # inheritance + >>> class DetDataSample(BaseDataElement): + ... @property + ... def proposals(self): + ... return self._proposals + ... @proposals.setter + ... def proposals(self, value): + ... self.set_field(value, '_proposals', dtype=BaseDataElement) + ... @proposals.deleter + ... def proposals(self): + ... del self._proposals + ... @property + ... def gt_instances(self): + ... return self._gt_instances + ... @gt_instances.setter + ... def gt_instances(self, value): + ... self.set_field(value, '_gt_instances', + ... dtype=BaseDataElement) + ... @gt_instances.deleter + ... def gt_instances(self): + ... del self._gt_instances + ... @property + ... def pred_instances(self): + ... return self._pred_instances + ... @pred_instances.setter + ... def pred_instances(self, value): + ... self.set_field(value, '_pred_instances', + ... dtype=BaseDataElement) + ... @pred_instances.deleter + ... def pred_instances(self): + ... del self._pred_instances + >>> det_sample = DetDataSample() + >>> proposals = BaseDataElement(bboxes=torch.rand((5, 4))) + >>> det_sample.proposals = proposals + >>> assert 'proposals' in det_sample + >>> assert det_sample.proposals == proposals + >>> del det_sample.proposals + >>> assert 'proposals' not in det_sample + >>> with self.assertRaises(AssertionError): + ... det_sample.proposals = torch.rand((5, 4)) + """ + + def __init__(self, *, metainfo: Optional[dict] = None, **kwargs) -> None: + self._metainfo_fields: set = set() + self._data_fields: set = set() + + if metainfo is not None: + self.set_metainfo(metainfo=metainfo) + if kwargs: + self.set_data(kwargs) + + def set_metainfo(self, metainfo: dict) -> None: + """Set or change key-value pairs in ``metainfo_field`` by parameter + ``metainfo``. + + Args: + metainfo (dict): A dict contains the meta information + of image, such as ``img_shape``, ``scale_factor``, etc. + """ + assert isinstance( + metainfo, dict + ), f"metainfo should be a ``dict`` but got {type(metainfo)}" + meta = copy.deepcopy(metainfo) + for k, v in meta.items(): + self.set_field(name=k, value=v, field_type="metainfo", dtype=None) + + def set_data(self, data: dict) -> None: + """Set or change key-value pairs in ``data_field`` by parameter + ``data``. + + Args: + data (dict): A dict contains annotations of image or + model predictions. + """ + assert isinstance(data, dict), f"data should be a `dict` but got {data}" + for k, v in data.items(): + # Use `setattr()` rather than `self.set_field` to allow `set_data` + # to set property method. + setattr(self, k, v) + + def update(self, instance: "BaseDataElement") -> None: + """The update() method updates the BaseDataElement with the elements + from another BaseDataElement object. + + Args: + instance (BaseDataElement): Another BaseDataElement object for + update the current object. + """ + assert isinstance( + instance, BaseDataElement + ), f"instance should be a `BaseDataElement` but got {type(instance)}" + self.set_metainfo(dict(instance.metainfo_items())) + self.set_data(dict(instance.items())) + + def new(self, *, metainfo: Optional[dict] = None, **kwargs) -> "BaseDataElement": + """Return a new data element with same type. If ``metainfo`` and + ``data`` are None, the new data element will have same metainfo and + data. If metainfo or data is not None, the new result will overwrite it + with the input value. + + Args: + metainfo (dict, optional): A dict contains the meta information + of image, such as ``img_shape``, ``scale_factor``, etc. + Defaults to None. + kwargs (dict): A dict contains annotations of image or + model predictions. + + Returns: + BaseDataElement: A new data element with same type. + """ + new_data = self.__class__() + + if metainfo is not None: + new_data.set_metainfo(metainfo) + else: + new_data.set_metainfo(dict(self.metainfo_items())) + if kwargs: + new_data.set_data(kwargs) + else: + new_data.set_data(dict(self.items())) + return new_data + + def clone(self): + """Deep copy the current data element. + + Returns: + BaseDataElement: The copy of current data element. + """ + clone_data = self.__class__() + clone_data.set_metainfo(dict(self.metainfo_items())) + clone_data.set_data(dict(self.items())) + return clone_data + + def keys(self) -> list: + """ + Returns: + list: Contains all keys in data_fields. + """ + # We assume that the name of the attribute related to property is + # '_' + the name of the property. We use this rule to filter out + # private keys. + # TODO: Use a more robust way to solve this problem + private_keys = { + "_" + key + for key in self._data_fields + if isinstance(getattr(type(self), key, None), property) + } + return list(self._data_fields - private_keys) + + def metainfo_keys(self) -> list: + """ + Returns: + list: Contains all keys in metainfo_fields. + """ + return list(self._metainfo_fields) + + def values(self) -> list: + """ + Returns: + list: Contains all values in data. + """ + return [getattr(self, k) for k in self.keys()] + + def metainfo_values(self) -> list: + """ + Returns: + list: Contains all values in metainfo. + """ + return [getattr(self, k) for k in self.metainfo_keys()] + + def all_keys(self) -> list: + """ + Returns: + list: Contains all keys in metainfo and data. + """ + return self.metainfo_keys() + self.keys() + + def all_values(self) -> list: + """ + Returns: + list: Contains all values in metainfo and data. + """ + return self.metainfo_values() + self.values() + + def all_items(self) -> Iterator[Tuple[str, Any]]: + """ + Returns: + iterator: An iterator object whose element is (key, value) tuple + pairs for ``metainfo`` and ``data``. + """ + for k in self.all_keys(): + yield (k, getattr(self, k)) + + def items(self) -> Iterator[Tuple[str, Any]]: + """ + Returns: + iterator: An iterator object whose element is (key, value) tuple + pairs for ``data``. + """ + for k in self.keys(): + yield (k, getattr(self, k)) + + def metainfo_items(self) -> Iterator[Tuple[str, Any]]: + """ + Returns: + iterator: An iterator object whose element is (key, value) tuple + pairs for ``metainfo``. + """ + for k in self.metainfo_keys(): + yield (k, getattr(self, k)) + + @property + def metainfo(self) -> dict: + """dict: A dict contains metainfo of current data element.""" + return dict(self.metainfo_items()) + + def __setattr__(self, name: str, value: Any): + """setattr is only used to set data.""" + if name in ("_metainfo_fields", "_data_fields"): + if not hasattr(self, name): + super().__setattr__(name, value) + else: + raise AttributeError( + f"{name} has been used as a " + "private attribute, which is immutable." + ) + else: + self.set_field(name=name, value=value, field_type="data", dtype=None) + + def __delattr__(self, item: str): + """Delete the item in dataelement. + + Args: + item (str): The key to delete. + """ + if item in ("_metainfo_fields", "_data_fields"): + raise AttributeError( + f"{item} has been used as a " "private attribute, which is immutable." + ) + super().__delattr__(item) + if item in self._metainfo_fields: + self._metainfo_fields.remove(item) + elif item in self._data_fields: + self._data_fields.remove(item) + + # dict-like methods + __delitem__ = __delattr__ + + def get(self, key, default=None) -> Any: + """Get property in data and metainfo as the same as python.""" + # Use `getattr()` rather than `self.__dict__.get()` to allow getting + # properties. + return getattr(self, key, default) + + def pop(self, *args) -> Any: + """Pop property in data and metainfo as the same as python.""" + assert len(args) < 3, "``pop`` get more than 2 arguments" + name = args[0] + if name in self._metainfo_fields: + self._metainfo_fields.remove(args[0]) + return self.__dict__.pop(*args) + + elif name in self._data_fields: + self._data_fields.remove(args[0]) + return self.__dict__.pop(*args) + + # with default value + elif len(args) == 2: + return args[1] + else: + # don't just use 'self.__dict__.pop(*args)' for only popping key in + # metainfo or data + raise KeyError(f"{args[0]} is not contained in metainfo or data") + + def __contains__(self, item: str) -> bool: + """Whether the item is in dataelement. + + Args: + item (str): The key to inquire. + """ + return item in self._data_fields or item in self._metainfo_fields + + def set_field( + self, + value: Any, + name: str, + dtype: Optional[Union[Type, Tuple[Type, ...]]] = None, + field_type: str = "data", + ) -> None: + """Special method for set union field, used as property.setter + functions.""" + assert field_type in ["metainfo", "data"] + if dtype is not None: + assert isinstance( + value, dtype + ), f"{value} should be a {dtype} but got {type(value)}" + + if field_type == "metainfo": + if name in self._data_fields: + raise AttributeError( + f"Cannot set {name} to be a field of metainfo " + f"because {name} is already a data field" + ) + self._metainfo_fields.add(name) + else: + if name in self._metainfo_fields: + raise AttributeError( + f"Cannot set {name} to be a field of data " + f"because {name} is already a metainfo field" + ) + self._data_fields.add(name) + super().__setattr__(name, value) + + # Tensor-like methods + def to(self, *args, **kwargs) -> "BaseDataElement": + """Apply same name function to all tensors in data_fields.""" + new_data = self.new() + for k, v in self.items(): + if hasattr(v, "to"): + v = v.to(*args, **kwargs) + data = {k: v} + new_data.set_data(data) + return new_data + + # Tensor-like methods + def cpu(self) -> "BaseDataElement": + """Convert all tensors to CPU in data.""" + new_data = self.new() + for k, v in self.items(): + if isinstance(v, (torch.Tensor, BaseDataElement)): + v = v.cpu() + data = {k: v} + new_data.set_data(data) + return new_data + + # Tensor-like methods + def cuda(self) -> "BaseDataElement": + """Convert all tensors to GPU in data.""" + new_data = self.new() + for k, v in self.items(): + if isinstance(v, (torch.Tensor, BaseDataElement)): + v = v.cuda() + data = {k: v} + new_data.set_data(data) + return new_data + + # Tensor-like methods + def npu(self) -> "BaseDataElement": + """Convert all tensors to NPU in data.""" + new_data = self.new() + for k, v in self.items(): + if isinstance(v, (torch.Tensor, BaseDataElement)): + v = v.npu() + data = {k: v} + new_data.set_data(data) + return new_data + + def mlu(self) -> "BaseDataElement": + """Convert all tensors to MLU in data.""" + new_data = self.new() + for k, v in self.items(): + if isinstance(v, (torch.Tensor, BaseDataElement)): + v = v.mlu() + data = {k: v} + new_data.set_data(data) + return new_data + + # Tensor-like methods + def detach(self) -> "BaseDataElement": + """Detach all tensors in data.""" + new_data = self.new() + for k, v in self.items(): + if isinstance(v, (torch.Tensor, BaseDataElement)): + v = v.detach() + data = {k: v} + new_data.set_data(data) + return new_data + + # Tensor-like methods + def numpy(self) -> "BaseDataElement": + """Convert all tensors to np.ndarray in data.""" + new_data = self.new() + for k, v in self.items(): + if isinstance(v, (torch.Tensor, BaseDataElement)): + v = v.detach().cpu().numpy() + data = {k: v} + new_data.set_data(data) + return new_data + + def to_tensor(self) -> "BaseDataElement": + """Convert all np.ndarray to tensor in data.""" + new_data = self.new() + for k, v in self.items(): + data = {} + if isinstance(v, np.ndarray): + v = torch.from_numpy(v) + data[k] = v + elif isinstance(v, BaseDataElement): + v = v.to_tensor() + data[k] = v + new_data.set_data(data) + return new_data + + def to_dict(self) -> dict: + """Convert BaseDataElement to dict.""" + return { + k: v.to_dict() if isinstance(v, BaseDataElement) else v + for k, v in self.all_items() + } + + def __repr__(self) -> str: + """Represent the object.""" + + def _addindent(s_: str, num_spaces: int) -> str: + """This func is modified from `pytorch` https://github.com/pytorch/ + pytorch/blob/b17b2b1cc7b017c3daaeff8cc7ec0f514d42ec37/torch/nn/modu + les/module.py#L29. + + Args: + s_ (str): The string to add spaces. + num_spaces (int): The num of space to add. + + Returns: + str: The string after add indent. + """ + s = s_.split("\n") + # don't do anything for single-line stuff + if len(s) == 1: + return s_ + first = s.pop(0) + s = [(num_spaces * " ") + line for line in s] + s = "\n".join(s) # type: ignore + s = first + "\n" + s # type: ignore + return s # type: ignore + + def dump(obj: Any) -> str: + """Represent the object. + + Args: + obj (Any): The obj to represent. + + Returns: + str: The represented str. + """ + _repr = "" + if isinstance(obj, dict): + for k, v in obj.items(): + _repr += f"\n{k}: {_addindent(dump(v), 4)}" + elif isinstance(obj, BaseDataElement): + _repr += "\n\n META INFORMATION" + metainfo_items = dict(obj.metainfo_items()) + _repr += _addindent(dump(metainfo_items), 4) + _repr += "\n\n DATA FIELDS" + items = dict(obj.items()) + _repr += _addindent(dump(items), 4) + classname = obj.__class__.__name__ + _repr = f"<{classname}({_repr}\n) at {hex(id(obj))}>" + else: + _repr += repr(obj) + return _repr + + return dump(self) diff --git a/abl/structures/list_data.py b/abl/structures/list_data.py new file mode 100644 index 0000000..2571a13 --- /dev/null +++ b/abl/structures/list_data.py @@ -0,0 +1,321 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import itertools +from collections.abc import Sized +from typing import Any, List, Union + +import numpy as np +import torch + +from ..utils import flatten as flatten_list +from ..utils import to_hashable +from .base_data_element import BaseDataElement + +BoolTypeTensor = Union[torch.BoolTensor, torch.cuda.BoolTensor] +LongTypeTensor = Union[torch.LongTensor, torch.cuda.LongTensor] + +IndexType = Union[str, slice, int, list, LongTypeTensor, BoolTypeTensor, np.ndarray] + + +# Modified from +# https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/data_structures/instance_data.py # noqa +class ListData(BaseDataElement): + """Data structure for instance-level annotations or predictions. + + Subclass of :class:`BaseDataElement`. All value in `data_fields` + should have the same length. This design refer to + https://github.com/facebookresearch/detectron2/blob/master/detectron2/structures/instances.py # noqa E501 + ListData also support extra functions: ``index``, ``slice`` and ``cat`` for data field. The type of value + in data field can be base data structure such as `torch.Tensor`, `numpy.ndarray`, `list`, `str`, `tuple`, + and can be customized data structure that has ``__len__``, ``__getitem__`` and ``cat`` attributes. + + Examples: + >>> # custom data structure + >>> class TmpObject: + ... def __init__(self, tmp) -> None: + ... assert isinstance(tmp, list) + ... self.tmp = tmp + ... def __len__(self): + ... return len(self.tmp) + ... def __getitem__(self, item): + ... if isinstance(item, int): + ... if item >= len(self) or item < -len(self): # type:ignore + ... raise IndexError(f'Index {item} out of range!') + ... else: + ... # keep the dimension + ... item = slice(item, None, len(self)) + ... return TmpObject(self.tmp[item]) + ... @staticmethod + ... def cat(tmp_objs): + ... assert all(isinstance(results, TmpObject) for results in tmp_objs) + ... if len(tmp_objs) == 1: + ... return tmp_objs[0] + ... tmp_list = [tmp_obj.tmp for tmp_obj in tmp_objs] + ... tmp_list = list(itertools.chain(*tmp_list)) + ... new_data = TmpObject(tmp_list) + ... return new_data + ... def __repr__(self): + ... return str(self.tmp) + >>> from mmengine.structures import ListData + >>> import numpy as np + >>> import torch + >>> img_meta = dict(img_shape=(800, 1196, 3), pad_shape=(800, 1216, 3)) + >>> instance_data = ListData(metainfo=img_meta) + >>> 'img_shape' in instance_data + True + >>> instance_data.det_labels = torch.LongTensor([2, 3]) + >>> instance_data["det_scores"] = torch.Tensor([0.8, 0.7]) + >>> instance_data.bboxes = torch.rand((2, 4)) + >>> instance_data.polygons = TmpObject([[1, 2, 3, 4], [5, 6, 7, 8]]) + >>> len(instance_data) + 2 + >>> print(instance_data) + + >>> sorted_results = instance_data[instance_data.det_scores.sort().indices] + >>> sorted_results.det_scores + tensor([0.7000, 0.8000]) + >>> print(instance_data[instance_data.det_scores > 0.75]) + + >>> print(instance_data[instance_data.det_scores > 1]) + + >>> print(instance_data.cat([instance_data, instance_data])) + + """ + + def __setattr__(self, name: str, value: list): + """setattr is only used to set data. + + The value must have the attribute of `__len__` and have the same length + of `ListData`. + """ + if name in ("_metainfo_fields", "_data_fields"): + if not hasattr(self, name): + super().__setattr__(name, value) + else: + raise AttributeError( + f"{name} has been used as a " + "private attribute, which is immutable." + ) + + else: + assert isinstance(value, list), "value must be of type `list`" + + if len(self) > 0: + assert len(value) == len(self), ( + "The length of " + f"values {len(value)} is " + "not consistent with " + "the length of this " + ":obj:`ListData` " + f"{len(self)}" + ) + super().__setattr__(name, value) + + __setitem__ = __setattr__ + + def __getitem__(self, item: IndexType) -> "ListData": + """ + Args: + item (str, int, list, :obj:`slice`, :obj:`numpy.ndarray`, + :obj:`torch.LongTensor`, :obj:`torch.BoolTensor`): + Get the corresponding values according to item. + + Returns: + :obj:`ListData`: Corresponding values. + """ + assert isinstance(item, IndexType.__args__) + if isinstance(item, list): + item = np.array(item) + if isinstance(item, np.ndarray): + # The default int type of numpy is platform dependent, int32 for + # windows and int64 for linux. `torch.Tensor` requires the index + # should be int64, therefore we simply convert it to int64 here. + # More details in https://github.com/numpy/numpy/issues/9464 + item = item.astype(np.int64) if item.dtype == np.int32 else item + item = torch.from_numpy(item) + + if isinstance(item, str): + return getattr(self, item) + + if isinstance(item, int): + if item >= len(self) or item < -len(self): # type:ignore + raise IndexError(f"Index {item} out of range!") + else: + # keep the dimension + item = slice(item, None, len(self)) + + new_data = self.__class__(metainfo=self.metainfo) + if isinstance(item, torch.Tensor): + assert item.dim() == 1, ( + "Only support to get the" " values along the first dimension." + ) + if isinstance(item, BoolTypeTensor.__args__): + assert len(item) == len(self), ( + "The shape of the " + "input(BoolTensor) " + f"{len(item)} " + "does not match the shape " + "of the indexed tensor " + "in results_field " + f"{len(self)} at " + "first dimension." + ) + + for k, v in self.items(): + if isinstance(v, torch.Tensor): + new_data[k] = v[item] + elif isinstance(v, np.ndarray): + new_data[k] = v[item.cpu().numpy()] + elif isinstance(v, (str, list, tuple)) or ( + hasattr(v, "__getitem__") and hasattr(v, "cat") + ): + # convert to indexes from BoolTensor + if isinstance(item, BoolTypeTensor.__args__): + indexes = torch.nonzero(item).view(-1).cpu().numpy().tolist() + else: + indexes = item.cpu().numpy().tolist() + slice_list = [] + if indexes: + for index in indexes: + slice_list.append(slice(index, None, len(v))) + else: + slice_list.append(slice(None, 0, None)) + r_list = [v[s] for s in slice_list] + if isinstance(v, (str, list, tuple)): + new_value = r_list[0] + for r in r_list[1:]: + new_value = new_value + r + else: + new_value = v.cat(r_list) + new_data[k] = new_value + else: + raise ValueError( + f"The type of `{k}` is `{type(v)}`, which has no " + "attribute of `cat`, so it does not " + "support slice with `bool`" + ) + + else: + # item is a slice + for k, v in self.items(): + new_data[k] = v[item] + return new_data # type:ignore + + @staticmethod + def cat(instances_list: List["ListData"]) -> "ListData": + """Concat the instances of all :obj:`ListData` in the list. + + Note: To ensure that cat returns as expected, make sure that + all elements in the list must have exactly the same keys. + + Args: + instances_list (list[:obj:`ListData`]): A list + of :obj:`ListData`. + + Returns: + :obj:`ListData` + """ + assert all(isinstance(results, ListData) for results in instances_list) + assert len(instances_list) > 0 + if len(instances_list) == 1: + return instances_list[0] + + # metainfo and data_fields must be exactly the + # same for each element to avoid exceptions. + field_keys_list = [instances.all_keys() for instances in instances_list] + assert len({len(field_keys) for field_keys in field_keys_list}) == 1 and len( + set(itertools.chain(*field_keys_list)) + ) == len(field_keys_list[0]), ( + "There are different keys in " + "`instances_list`, which may " + "cause the cat operation " + "to fail. Please make sure all " + "elements in `instances_list` " + "have the exact same key." + ) + + new_data = instances_list[0].__class__(metainfo=instances_list[0].metainfo) + for k in instances_list[0].keys(): + values = [results[k] for results in instances_list] + v0 = values[0] + if isinstance(v0, torch.Tensor): + new_values = torch.cat(values, dim=0) + elif isinstance(v0, np.ndarray): + new_values = np.concatenate(values, axis=0) + elif isinstance(v0, (str, list, tuple)): + new_values = v0[:] + for v in values[1:]: + new_values += v + elif hasattr(v0, "cat"): + new_values = v0.cat(values) + else: + raise ValueError( + f"The type of `{k}` is `{type(v0)}` which has no " + "attribute of `cat`" + ) + new_data[k] = new_values + return new_data # type:ignore + + def flatten(self, item: IndexType) -> List: + """Flatten self[item]. + + Returns: + list: Flattened data fields. + """ + return flatten_list(self[item]) + + def elements_num(self, item: IndexType) -> int: + """int: The number of elements in self[item].""" + return len(self.flatten(item)) + + def to_tuple(self, item: IndexType) -> tuple: + """tuple: The data fields in self[item] converted to tuple.""" + return to_hashable(self[item]) + + def __len__(self) -> int: + """int: The length of ListData.""" + if len(self._data_fields) > 0: + one_element = next(iter(self._data_fields)) + return len(getattr(self, one_element)) + # return len(self.values()[0]) + else: + return 0 diff --git a/abl/utils/__init__.py b/abl/utils/__init__.py index 75c7990..526b50b 100644 --- a/abl/utils/__init__.py +++ b/abl/utils/__init__.py @@ -1,2 +1,3 @@ +from .cache import Cache from .logger import ABLLogger, print_log -from .utils import * \ No newline at end of file +from .utils import * diff --git a/abl/utils/cache.py b/abl/utils/cache.py new file mode 100644 index 0000000..f4b3b0c --- /dev/null +++ b/abl/utils/cache.py @@ -0,0 +1,112 @@ +import pickle +from os import PathLike +from pathlib import Path +from typing import Callable, Generic, Hashable, TypeVar, Union + +from .logger import print_log + +K = TypeVar("K") +T = TypeVar("T") +PREV, NEXT, KEY, RESULT = 0, 1, 2, 3 # names for the link fields + + +class Cache(Generic[K, T]): + def __init__( + self, + func: Callable[[K], T], + cache: bool, + cache_file: Union[None, str, PathLike], + key_func: Callable[[K], Hashable] = lambda x: x, + max_size: int = 4096, + ): + """Create cache + + :param func: Function this cache evaluates + :param cache: If true, do in memory caching. + :param cache_root: If not None, cache to files at the provided path. + :param key_func: Convert the key into a hashable object if needed + """ + self.func = func + self.key_func = key_func + self.cache = cache + if cache is True or cache_file is not None: + print_log("Caching is activated", logger="current") + self._init_cache(cache_file, max_size) + self.first = self.get_from_dict + else: + self.first = self.func + + def __getitem__(self, item: K, *args) -> T: + return self.first(item, *args) + + def invalidate(self): + """Invalidate entire cache.""" + self.cache_dict.clear() + if self.cache_file: + for p in self.cache_root.iterdir(): + p.unlink() + + def _init_cache(self, cache_file, max_size): + self.cache = True + self.cache_dict = dict() + + self.hits, self.misses, self.maxsize = 0, 0, max_size + self.full = False + self.root = [] # root of the circular doubly linked list + self.root[:] = [self.root, self.root, None, None] + + if cache_file is not None: + with open(cache_file, "rb") as f: + cache_dict_from_file = pickle.load(f) + self.maxsize += len(cache_dict_from_file) + print_log( + f"Max size of the cache has been enlarged to {self.maxsize}.", logger="current" + ) + for cache_key, result in cache_dict_from_file.items(): + last = self.root[PREV] + link = [last, self.root, cache_key, result] + last[NEXT] = self.root[PREV] = self.cache_dict[cache_key] = link + + def get(self, item: K, *args) -> T: + return self.first(item, *args) + + def get_from_dict(self, item: K, *args) -> T: + """Implements dict based cache.""" + cache_key = (self.key_func(item), *args) + link = self.cache_dict.get(cache_key) + if link is not None: + # Move the link to the front of the circular queue + link_prev, link_next, _key, result = link + link_prev[NEXT] = link_next + link_next[PREV] = link_prev + last = self.root[PREV] + last[NEXT] = self.root[PREV] = link + link[PREV] = last + link[NEXT] = self.root + self.hits += 1 + return result + self.misses += 1 + + result = self.func(item, *args) + + if self.full: + # Use the old root to store the new key and result. + oldroot = self.root + oldroot[KEY] = cache_key + oldroot[RESULT] = result + # Empty the oldest link and make it the new root. + self.root = oldroot[NEXT] + oldkey = self.root[KEY] + oldresult = self.root[RESULT] + self.root[KEY] = self.root[RESULT] = None + # Now update the cache dictionary. + del self.cache_dict[oldkey] + self.cache_dict[cache_key] = oldroot + else: + # Put result in a new link at the front of the queue. + last = self.root[PREV] + link = [last, self.root, cache_key, result] + last[NEXT] = self.root[PREV] = self.cache_dict[cache_key] = link + if isinstance(self.maxsize, int): + self.full = len(self.cache_dict) >= self.maxsize + return result diff --git a/abl/utils/utils.py b/abl/utils/utils.py index 485cf1a..8192bf9 100644 --- a/abl/utils/utils.py +++ b/abl/utils/utils.py @@ -1,6 +1,7 @@ -import numpy as np from itertools import chain +import numpy as np + def flatten(nested_list): """ @@ -15,6 +16,11 @@ def flatten(nested_list): ------- list A flattened version of the input list. + + Raises + ------ + TypeError + If the input object is not a list. """ if not isinstance(nested_list, list): raise TypeError("Input must be of type list.") @@ -41,6 +47,9 @@ def reform_idx(flattened_list, structured_list): list A reformed list that mimics the structure of structured_list. """ + # if not isinstance(flattened_list, list): + # raise TypeError("Input must be of type list.") + if not isinstance(structured_list[0], (list, tuple)): return flattened_list @@ -80,7 +89,7 @@ def hamming_dist(pred_pseudo_label, candidates): return np.sum(pred_pseudo_label != candidates, axis=1) -def confidence_dist(pred_prob, candidates_idx): +def confidence_dist(pred_prob, candidates): """ Compute the confidence distance between prediction probabilities and candidates. @@ -89,7 +98,7 @@ def confidence_dist(pred_prob, candidates_idx): pred_prob : list of numpy.ndarray Prediction probability distributions, each element is an ndarray representing the probability distribution of a particular prediction. - candidates_idx : list of list of int + candidates : list of list of int Index of candidate labels, each element is a list of indexes being considered as a candidate correction. @@ -99,8 +108,8 @@ def confidence_dist(pred_prob, candidates_idx): Confidence distances computed for each candidate. """ pred_prob = np.clip(pred_prob, 1e-9, 1) - _, cols = np.indices((len(candidates_idx), len(candidates_idx[0]))) - return 1 - np.prod(pred_prob[cols, candidates_idx], axis=1) + _, cols = np.indices((len(candidates), len(candidates[0]))) + return 1 - np.prod(pred_prob[cols, candidates], axis=1) def block_sample(X, Z, Y, sample_num, seg_idx): @@ -135,6 +144,34 @@ def block_sample(X, Z, Y, sample_num, seg_idx): return (data[start_idx:end_idx] for data in (X, Z, Y)) +def check_equal(a, b, max_err=0): + """ + Check whether two numbers a and b are equal within a maximum allowable error. + + Parameters + ---------- + a, b : int or float + The numbers to compare. + max_err : int or float, optional + The maximum allowable absolute difference between a and b for them to be considered equal. + Default is 0, meaning the numbers must be exactly equal. + + Returns + ------- + bool + True if a and b are equal within the allowable error, False otherwise. + + Raises + ------ + TypeError + If a or b are not of type int or float. + """ + if not (isinstance(a, (int, float)) and isinstance(b, (int, float))): + raise TypeError("Input values must be int or float.") + + return abs(a - b) <= max_err + + def to_hashable(x): """ Convert a nested list to a nested tuple so it is hashable. @@ -154,7 +191,8 @@ def to_hashable(x): return tuple(to_hashable(item) for item in x) return x -def restore_from_hashable(x): + +def hashable_to_list(x): """ Convert a nested tuple back to a nested list. @@ -174,6 +212,45 @@ def restore_from_hashable(x): return x +def calculate_revision_num(parameter, total_length): + """ + Convert a float parameter to an integer, based on a total length. + + Parameters + ---------- + parameter : int or float + The parameter to convert. If float, it should be between 0 and 1. + If int, it should be non-negative. If -1, it will be replaced with total_length. + total_length : int + The total length to calculate the parameter from if it's a fraction. + + Returns + ------- + int + The calculated parameter. + + Raises + ------ + TypeError + If parameter is not an int or a float. + ValueError + If parameter is a float not in [0, 1] or an int below 0. + """ + if not isinstance(parameter, (int, float)): + raise TypeError("Parameter must be of type int or float.") + + if parameter == -1: + return total_length + elif isinstance(parameter, float): + if not (0 <= parameter <= 1): + raise ValueError("If parameter is a float, it must be between 0 and 1.") + return round(total_length * parameter) + else: + if parameter < 0: + raise ValueError("If parameter is an int, it must be non-negative.") + return parameter + + if __name__ == "__main__": A = np.array( [ @@ -227,4 +304,5 @@ if __name__ == "__main__": ) B = [[0, 9, 3], [0, 11, 4]] + print(ori_confidence_dist(A, B)) print(confidence_dist(A, B))