/
_tensorboard_logger.py
179 lines (151 loc) · 7.29 KB
/
_tensorboard_logger.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
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
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# 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.
"""Contains a logger to push training logs to the Hub, using Tensorboard."""
from pathlib import Path
from typing import TYPE_CHECKING, List, Optional, Union
from huggingface_hub._commit_scheduler import CommitScheduler
from .utils import experimental
# Depending on user's setup, SummaryWriter can come either from 'tensorboardX'
# or from 'torch.utils.tensorboard'. Both are compatible so let's try to load
# from either of them.
try:
from tensorboardX import SummaryWriter
is_summary_writer_available = True
except ImportError:
try:
from torch.utils.tensorboard import SummaryWriter
is_summary_writer_available = False
except ImportError:
# Dummy class to avoid failing at import. Will raise on instance creation.
SummaryWriter = object
is_summary_writer_available = False
if TYPE_CHECKING:
from tensorboardX import SummaryWriter
class HFSummaryWriter(SummaryWriter):
"""
Wrapper around the tensorboard's `SummaryWriter` to push training logs to the Hub.
Data is logged locally and then pushed to the Hub asynchronously. Pushing data to the Hub is done in a separate
thread to avoid blocking the training script. In particular, if the upload fails for any reason (e.g. a connection
issue), the main script will not be interrupted. Data is automatically pushed to the Hub every `commit_every`
minutes (default to every 5 minutes).
<Tip warning={true}>
`HFSummaryWriter` is experimental. Its API is subject to change in the future without prior notice.
</Tip>
Args:
repo_id (`str`):
The id of the repo to which the logs will be pushed.
logdir (`str`, *optional*):
The directory where the logs will be written. If not specified, a local directory will be created by the
underlying `SummaryWriter` object.
commit_every (`int` or `float`, *optional*):
The frequency (in minutes) at which the logs will be pushed to the Hub. Defaults to 5 minutes.
squash_history (`bool`, *optional*):
Whether to squash the history of the repo after each commit. Defaults to `False`. Squashing commits is
useful to avoid degraded performances on the repo when it grows too large.
repo_type (`str`, *optional*):
The type of the repo to which the logs will be pushed. Defaults to "model".
repo_revision (`str`, *optional*):
The revision of the repo to which the logs will be pushed. Defaults to "main".
repo_private (`bool`, *optional*):
Whether to create a private repo or not. Defaults to False. This argument is ignored if the repo already
exists.
path_in_repo (`str`, *optional*):
The path to the folder in the repo where the logs will be pushed. Defaults to "tensorboard/".
repo_allow_patterns (`List[str]` or `str`, *optional*):
A list of patterns to include in the upload. Defaults to `"*.tfevents.*"`. Check out the
[upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details.
repo_ignore_patterns (`List[str]` or `str`, *optional*):
A list of patterns to exclude in the upload. Check out the
[upload guide](https://huggingface.co/docs/huggingface_hub/guides/upload#upload-a-folder) for more details.
token (`str`, *optional*):
Authentication token. Will default to the stored token. See https://huggingface.co/settings/token for more
details
kwargs:
Additional keyword arguments passed to `SummaryWriter`.
Examples:
```py
>>> from huggingface_hub import HFSummaryWriter
# Logs are automatically pushed every 15 minutes
>>> logger = HFSummaryWriter(repo_id="test_hf_logger", commit_every=15)
>>> logger.add_scalar("a", 1)
>>> logger.add_scalar("b", 2)
...
# You can also trigger a push manually
>>> logger.scheduler.trigger()
```
```py
>>> from huggingface_hub import HFSummaryWriter
# Logs are automatically pushed every 5 minutes (default) + when exiting the context manager
>>> with HFSummaryWriter(repo_id="test_hf_logger") as logger:
... logger.add_scalar("a", 1)
... logger.add_scalar("b", 2)
```
"""
@experimental
def __new__(cls, *args, **kwargs) -> "HFSummaryWriter":
if not is_summary_writer_available:
raise ImportError(
"You must have `tensorboard` installed to use `HFSummaryWriter`. Please run `pip install --upgrade"
" tensorboardX` first."
)
return super().__new__(cls)
def __init__(
self,
repo_id: str,
*,
logdir: Optional[str] = None,
commit_every: Union[int, float] = 5,
squash_history: bool = False,
repo_type: Optional[str] = None,
repo_revision: Optional[str] = None,
repo_private: bool = False,
path_in_repo: Optional[str] = "tensorboard",
repo_allow_patterns: Optional[Union[List[str], str]] = "*.tfevents.*",
repo_ignore_patterns: Optional[Union[List[str], str]] = None,
token: Optional[str] = None,
**kwargs,
):
# Initialize SummaryWriter
super().__init__(logdir=logdir, **kwargs)
# Check logdir has been correctly initialized and fail early otherwise. In practice, SummaryWriter takes care of it.
if not isinstance(self.logdir, str):
raise ValueError(f"`self.logdir` must be a string. Got '{self.logdir}' of type {type(self.logdir)}.")
# Append logdir name to `path_in_repo`
if path_in_repo is None or path_in_repo == "":
path_in_repo = Path(self.logdir).name
else:
path_in_repo = path_in_repo.strip("/") + "/" + Path(self.logdir).name
# Initialize scheduler
self.scheduler = CommitScheduler(
folder_path=self.logdir,
path_in_repo=path_in_repo,
repo_id=repo_id,
repo_type=repo_type,
revision=repo_revision,
private=repo_private,
token=token,
allow_patterns=repo_allow_patterns,
ignore_patterns=repo_ignore_patterns,
every=commit_every,
squash_history=squash_history,
)
# Exposing some high-level info at root level
self.repo_id = self.scheduler.repo_id
self.repo_type = self.scheduler.repo_type
self.repo_revision = self.scheduler.revision
def __exit__(self, exc_type, exc_val, exc_tb):
"""Push to hub in a non-blocking way when exiting the logger's context manager."""
super().__exit__(exc_type, exc_val, exc_tb)
future = self.scheduler.trigger()
future.result()