output.py (2065B)
1 from typing import Any, Dict, Optional, Type 2 import pathlib 3 import pickle 4 import types 5 6 import torch 7 8 import mmsrl.utils 9 10 11 class Output: 12 """ 13 Class for outputing data about the learned model. 14 """ 15 # TODO parametrize this class 16 17 def __init__(self, config: mmsrl.utils.dotdict, split: str, epoch: int, final: bool) -> None: 18 """ Initialize the outputs variables, but do not acquire necessary resources yet. """ 19 self.config: mmsrl.utils.dotdict = config 20 self.split = {"unseen_test": "test"}.get(split, split) 21 self.epoch = epoch 22 self.final = final 23 self.path = None 24 self.file = None 25 26 def __enter__(self): 27 """ Acquire resources needed for outputing the data. """ 28 if self.final and self.config.get(f"output_{self.split}"): 29 self.path = pathlib.Path(self.config[f"output_{self.split}"]) 30 if '{' in str(self.path): 31 self.path = pathlib.Path(str(path).format(self.epoch)) 32 self.file = self.path.open("wb") 33 return self 34 35 def __exit__(self, 36 exc_type: Optional[Type[BaseException]], 37 exc_inst: Optional[BaseException], 38 exc_tb: Optional[types.TracebackType]) -> None: 39 """ Free used resources. """ 40 if self.file is not None: 41 self.file.close() 42 43 def __call__(self, batch: Dict[str, Any], predictions: Optional[torch.Tensor]) -> None: 44 """ 45 Update model output data files with the given batch and the outputs of the model on this batch. 46 47 Args: 48 batch: the input values used for evaluation 49 predictions: the output of the model 50 """ 51 if self.file and predictions is not None: 52 for ib in range(len(batch["image_name"])): 53 entities = dict(zip( 54 filter(None, batch["entities_name"][ib]), 55 predictions[ib].detach().cpu())) 56 pickle.dump((batch["image_name"][ib], entities), self.file) 57 self.file.flush()