Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
ColossalAI
Commits
8c66a1d0
Unverified
Commit
8c66a1d0
authored
Nov 16, 2022
by
Jiarui Fang
Committed by
GitHub
Nov 16, 2022
Browse files
[polish] remove useless file _mem_tracer_hook.py (#1963)
parent
c4739a72
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
0 additions
and
44 deletions
+0
-44
colossalai/trainer/hooks/_mem_tracer_hook.py
colossalai/trainer/hooks/_mem_tracer_hook.py
+0
-44
No files found.
colossalai/trainer/hooks/_mem_tracer_hook.py
deleted
100644 → 0
View file @
c4739a72
from
colossalai.registry
import
HOOKS
from
torch
import
Tensor
from
colossalai.trainer.hooks
import
BaseHook
from
colossalai.gemini.memory_tracer
import
AsyncMemoryMonitor
@
HOOKS
.
register_module
class
MemTraceHook
(
BaseHook
):
"""Save memory stats and pass it to states
This hook is used to record memory usage info, and pass to trainer.states
You can use it as other trainer hook and fetch data from trainer.states['metrics][mode]
"""
def
__init__
(
self
,
priority
:
int
=
0
,
)
->
None
:
super
().
__init__
(
priority
=
priority
)
self
.
_memory_monitor
=
AsyncMemoryMonitor
()
def
after_hook_is_attached
(
self
,
trainer
):
# Initialize the data
trainer
.
states
[
'metrics'
][
'train'
]
=
self
.
_memory_monitor
.
state_dict
trainer
.
states
[
'metrics'
][
'test'
]
=
self
.
_memory_monitor
.
state_dict
def
before_train_iter
(
self
,
trainer
):
self
.
_memory_monitor
.
start
()
return
super
().
before_train_iter
(
trainer
)
def
after_train_iter
(
self
,
trainer
,
output
:
Tensor
,
label
:
Tensor
,
loss
:
Tensor
):
self
.
_memory_monitor
.
finish
()
trainer
.
states
[
'metrics'
][
'train'
]
=
self
.
_memory_monitor
.
state_dict
trainer
.
states
[
'metrics'
][
'test'
]
=
self
.
_memory_monitor
.
state_dict
return
super
().
after_train_iter
(
trainer
,
output
,
label
,
loss
)
def
before_test_iter
(
self
,
trainer
):
self
.
_memory_monitor
.
start
()
return
super
().
before_test
(
trainer
)
def
after_test_iter
(
self
,
trainer
,
output
:
Tensor
,
label
:
Tensor
,
loss
:
Tensor
):
self
.
_memory_monitor
.
finish
()
trainer
.
states
[
'metrics'
][
'train'
]
=
self
.
_memory_monitor
.
state_dict
trainer
.
states
[
'metrics'
][
'test'
]
=
self
.
_memory_monitor
.
state_dict
return
super
().
after_test_iter
(
trainer
,
output
,
label
,
loss
)
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment