-
Notifications
You must be signed in to change notification settings - Fork 1.1k
/
garbage_collector.py
87 lines (75 loc) · 3.54 KB
/
garbage_collector.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
# 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
import gc
from typing import TYPE_CHECKING
from monai.utils import IgniteInfo, min_version, optional_import
if TYPE_CHECKING:
from ignite.engine import Engine, Events
from ignite.engine.events import CallableEventWithFilter
else:
CallableEventWithFilter, _ = optional_import(
"ignite.engine.events", IgniteInfo.OPT_IMPORT_VERSION, min_version, "CallableEventWithFilter"
)
Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine")
Events, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Events")
class GarbageCollector:
"""
Run garbage collector after each epoch
Args:
trigger_event: the event that trigger a call to this handler.
- "epoch", after completion of each epoch (equivalent of ignite.engine.Events.EPOCH_COMPLETED)
- "iteration", after completion of each iteration (equivalent of ignite.engine.Events.ITERATION_COMPLETED)
- any ignite built-in event from ignite.engine.Events.
Defaults to "epoch".
log_level: log level (integer) for some garbage collection information as below. Defaults to 10 (DEBUG).
- 50 (CRITICAL)
- 40 (ERROR)
- 30 (WARNING)
- 20 (INFO)
- 10 (DEBUG)
- 0 (NOTSET)
"""
def __init__(self, trigger_event: str | Events | CallableEventWithFilter = "epoch", log_level: int = 10):
self.trigger_event: Events | CallableEventWithFilter
if isinstance(trigger_event, (Events, CallableEventWithFilter)):
self.trigger_event = trigger_event
elif trigger_event.lower() == "epoch":
self.trigger_event = Events.EPOCH_COMPLETED
elif trigger_event.lower() == "iteration":
self.trigger_event = Events.ITERATION_COMPLETED
else:
raise ValueError(
f"'trigger_event' should be either epoch, iteration, or an ignite built-in event from"
f" ignite.engine.Events, '{trigger_event}' was given."
)
self.log_level = log_level
def attach(self, engine: Engine) -> None:
if not engine.has_event_handler(self, self.trigger_event):
engine.add_event_handler(self.trigger_event, self)
def __call__(self, engine: Engine) -> None:
"""
This method calls python garbage collector.
Args:
engine: Ignite Engine, it should be either a trainer or validator.
"""
# get count before garbage collection
pre_count = gc.get_count()
# first call to garbage collector
gc.collect()
# second call to garbage collector
unreachable = gc.collect()
# get count after garbage collection
after_count = gc.get_count()
engine.logger.log(
self.log_level,
f"Garbage Count: [before: {pre_count}] -> [after: {after_count}] (unreachable : {unreachable})",
)