Source code for monai.handlers.smartcache_handler
# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
from typing import TYPE_CHECKING
from monai.config import IgniteInfo
from monai.data import SmartCacheDataset
from monai.utils import min_version, optional_import
Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events")
if TYPE_CHECKING:
from ignite.engine import Engine
else:
Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine")
[docs]
class SmartCacheHandler:
"""
Attach SmartCache logic to the engine in Ignite.
Mainly include the `start`, `update_cache`, and `shutdown` functions of SmartCacheDataset.
"""
[docs]
def __init__(self, smartcacher: SmartCacheDataset) -> None:
"""
Args:
smartcacher: predefined SmartCacheDataset, will attach it to the engine.
Raises:
TypeError: When ``smartcacher`` is not a ``monai.data.SmartCacheDataset``.
"""
if not isinstance(smartcacher, SmartCacheDataset):
raise TypeError("smartcacher must be a monai.data.SmartCacheDataset.")
self.smartcacher = smartcacher
[docs]
def attach(self, engine: Engine) -> None:
"""
Args:
engine: Ignite Engine, it can be a trainer, validator or evaluator.
"""
engine.add_event_handler(Events.STARTED, self.started)
engine.add_event_handler(Events.EPOCH_COMPLETED, self.epoch_completed)
engine.add_event_handler(Events.COMPLETED, self.completed)
[docs]
def started(self, engine: Engine) -> None:
"""Callback for train or validation/evaluation started Event.
Start the replacement thread of SmartCacheDataset.
Args:
engine: Ignite Engine, it can be a trainer, validator or evaluator.
"""
self.smartcacher.start()
[docs]
def epoch_completed(self, engine: Engine) -> None:
"""Callback for train or validation/evaluation epoch completed Event.
Update cache content with replacement data.
Args:
engine: Ignite Engine, it can be a trainer, validator or evaluator.
"""
self.smartcacher.update_cache()
[docs]
def completed(self, engine: Engine) -> None:
"""Callback for train or validation/evaluation completed Event.
Stop the replacement thread of SmartCacheDataset.
Args:
engine: Ignite Engine, it can be a trainer, validator or evaluator.
"""
self.smartcacher.shutdown()