forked from zycv/speechbrain
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtrain_logger.py
More file actions
202 lines (172 loc) · 6.52 KB
/
Copy pathtrain_logger.py
File metadata and controls
202 lines (172 loc) · 6.52 KB
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
"""Loggers for experiment monitoring.
Authors
* Peter Plantinga 2020
"""
import logging
import ruamel.yaml
logger = logging.getLogger(__name__)
class TrainLogger:
"""Abstract class defining an interface for training loggers."""
def log_stats(
self,
stats_meta,
train_stats=None,
valid_stats=None,
test_stats=None,
verbose=False,
):
"""Log the stats for one epoch.
Arguments
---------
stats_meta : dict of str:scalar pairs
Meta information about the stats (e.g., epoch, learning-rate, etc.).
train_stats : dict of str:list pairs
Each loss type is represented with a str : list pair including
all the values for the training pass.
valid_stats : dict of str:list pairs
Each loss type is represented with a str : list pair including
all the values for the validation pass.
test_stats : dict of str:list pairs
Each loss type is represented with a str : list pair including
all the values for the test pass.
verbose : bool
Whether to also put logging information to the standard logger.
"""
raise NotImplementedError
class FileTrainLogger(TrainLogger):
"""Text logger of training information.
Arguments
---------
save_file : str
The file to use for logging train information.
precision : int
Number of decimal places to display. Default 2, example: 1.35e-5.
summary_fns : dict of str:function pairs
Each summary function should take a list produced as output
from a training/validation pass and summarize it to a single scalar.
"""
def __init__(self, save_file, precision=2):
self.save_file = save_file
self.precision = precision
def _item_to_string(self, key, value, dataset=None):
"""Convert one item to string, handling floats"""
if isinstance(value, float) and 1.0 < value < 100.0:
value = f"{value:.{self.precision}f}"
elif isinstance(value, float):
value = f"{value:.{self.precision}e}"
if dataset is not None:
key = f"{dataset} {key}"
return f"{key}: {value}"
def _stats_to_string(self, stats, dataset=None):
"""Convert all stats to a single string summary"""
return ", ".join(
[self._item_to_string(k, v, dataset) for k, v in stats.items()]
)
def log_stats(
self,
stats_meta,
train_stats=None,
valid_stats=None,
test_stats=None,
verbose=True,
):
"""See TrainLogger.log_stats()"""
string_summary = self._stats_to_string(stats_meta)
for dataset, stats in [
("train", train_stats),
("valid", valid_stats),
("test", test_stats),
]:
if stats is not None:
string_summary += " - " + self._stats_to_string(stats, dataset)
with open(self.save_file, "a") as fout:
print(string_summary, file=fout)
if verbose:
logger.info(string_summary)
class TensorboardLogger(TrainLogger):
"""Logs training information in the format required by Tensorboard.
Arguments
---------
save_dir : str
A directory for storing all the relevant logs.
Raises
------
ImportError if Tensorboard is not installed.
"""
def __init__(self, save_dir):
self.save_dir = save_dir
# Raises ImportError if TensorBoard is not installed
from torch.utils.tensorboard import SummaryWriter
self.writer = SummaryWriter(self.save_dir)
self.global_step = {"train": {}, "valid": {}, "test": {}, "meta": 0}
def log_stats(
self,
stats_meta,
train_stats=None,
valid_stats=None,
test_stats=None,
verbose=False,
):
"""See TrainLogger.log_stats()"""
self.global_step["meta"] += 1
for name, value in stats_meta.items():
self.writer.add_scalar(name, value, self.global_step["meta"])
for dataset, stats in [
("train", train_stats),
("valid", valid_stats),
("test", test_stats),
]:
if stats is None:
continue
for stat, value_list in stats.items():
if stat not in self.global_step[dataset]:
self.global_step[dataset][stat] = 0
tag = f"{stat}/{dataset}"
# Both single value (per Epoch) and list (Per batch) logging is supported
if isinstance(value_list, list):
for value in value_list:
new_global_step = self.global_step[dataset][stat] + 1
self.writer.add_scalar(tag, value, new_global_step)
self.global_step[dataset][stat] = new_global_step
else:
value = value_list
new_global_step = self.global_step[dataset][stat] + 1
self.writer.add_scalar(tag, value, new_global_step)
self.global_step[dataset][stat] = new_global_step
class WandBLogger(TrainLogger):
"""Logger for wandb. To be used the same way as TrainLogger. Handles nested dicts as well.
An example on how to use this can be found in recipes/Voicebank/MTL/CoopNet/"""
def __init__(self, *args, **kwargs):
try:
yaml_file = kwargs.pop("yaml_config")
with open(yaml_file, "r") as yaml_stream:
# Read yaml with ruamel to ignore bangs
config_dict = ruamel.yaml.YAML().load(yaml_stream)
self.run = kwargs.pop("initializer", None)(
*args, **kwargs, config=config_dict
)
except Exception as e:
raise e("There was an issue with the WandB Logger initialization")
def log_stats(
self,
stats_meta,
train_stats=None,
valid_stats=None,
test_stats=None,
verbose=False,
):
"""See TrainLogger.log_stats()"""
logs = {}
for dataset, stats in [
("train", train_stats),
("valid", valid_stats),
("test", test_stats),
]:
if stats is None:
continue
logs[dataset] = stats
step = stats_meta.get("epoch", None)
if step is not None: # Useful for continuing runs that crashed
self.run.log({**logs, **stats_meta}, step=step)
else:
self.run.log({**logs, **stats_meta})