cfskit 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
cfskit-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 NUIST-GenAI-Lab
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
cfskit-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,44 @@
1
+ Metadata-Version: 2.4
2
+ Name: cfskit
3
+ Version: 0.1.0
4
+ Summary: Common research utilities for our NUIST-GenAI-Lab and CCGM tasks
5
+ Author-email: CFuShn <1354809038@qq.com>
6
+ License-Expression: MIT
7
+ Requires-Python: >=3.8
8
+ Description-Content-Type: text/markdown
9
+ License-File: LICENSE
10
+ Requires-Dist: loguru<1,>=0.6.0
11
+ Requires-Dist: accelerate<1,>=0.23.0; python_version < "3.10"
12
+ Requires-Dist: accelerate<2,>=1.0.0; python_version >= "3.10"
13
+ Requires-Dist: tensorboard<2.20,>=2.12.0; python_version < "3.9"
14
+ Requires-Dist: tensorboard<3,>=2.20.0; python_version >= "3.9"
15
+ Provides-Extra: torch
16
+ Requires-Dist: torch<3,>=2.0.0; extra == "torch"
17
+ Provides-Extra: tb
18
+ Requires-Dist: torch<3,>=2.0.0; extra == "tb"
19
+ Provides-Extra: tb-vision
20
+ Requires-Dist: torchvision<1,>=0.15.0; extra == "tb-vision"
21
+ Dynamic: license-file
22
+
23
+ # cfskit
24
+ Common research utilities for our NUIST-GenAI-Lab and CCGM tasks
25
+
26
+
27
+
28
+ ## Features
29
+
30
+ - 整合常用的科研库, 减少工具学习成本
31
+ - 尽量规范
32
+ - 尽量简洁
33
+ - 约定大于配置
34
+ - 最小化上手难度
35
+
36
+
37
+
38
+ ## 整合功能
39
+
40
+ - accelerate封装
41
+ - tensorboard封装
42
+ - loguru封装
43
+ - ipc 信号中断工具
44
+ - 常用训练,评估,测试工具
cfskit-0.1.0/README.md ADDED
@@ -0,0 +1,22 @@
1
+ # cfskit
2
+ Common research utilities for our NUIST-GenAI-Lab and CCGM tasks
3
+
4
+
5
+
6
+ ## Features
7
+
8
+ - 整合常用的科研库, 减少工具学习成本
9
+ - 尽量规范
10
+ - 尽量简洁
11
+ - 约定大于配置
12
+ - 最小化上手难度
13
+
14
+
15
+
16
+ ## 整合功能
17
+
18
+ - accelerate封装
19
+ - tensorboard封装
20
+ - loguru封装
21
+ - ipc 信号中断工具
22
+ - 常用训练,评估,测试工具
@@ -0,0 +1,39 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "cfskit"
7
+ version = "0.1.0"
8
+ description = "Common research utilities for our NUIST-GenAI-Lab and CCGM tasks"
9
+ authors = [{name="CFuShn", email="1354809038@qq.com"}]
10
+ license = "MIT"
11
+ readme = "README.md"
12
+ requires-python = ">=3.8"
13
+ dependencies = [
14
+ # 默认安装:轻量但常用的科研工具链
15
+ "loguru>=0.6.0,<1",
16
+ # accelerate 新版本往往要求 Python>=3.10,所以这里按 Python 版本分段
17
+ "accelerate>=0.23.0,<1; python_version < '3.10'",
18
+ "accelerate>=1.0.0,<2; python_version >= '3.10'",
19
+ # tensorboard 2.20 起需要 Python>=3.9,按 Python 版本分段
20
+ "tensorboard>=2.12.0,<2.20; python_version < '3.9'",
21
+ "tensorboard>=2.20.0,<3; python_version >= '3.9'",
22
+ ]
23
+
24
+ [project.optional-dependencies]
25
+ # 可选:大件依赖按需装
26
+ torch = ["torch>=2.0.0,<3"]
27
+ # 兼容原来的命名习惯:tb 只负责补 torch(tensorboard 已默认安装)
28
+ tb = ["torch>=2.0.0,<3"]
29
+ tb-vision = [
30
+ # 仅用于 torchvision.utils.make_grid(image_grid 等)
31
+ "torchvision>=0.15.0,<1",
32
+ ]
33
+
34
+ [tool.setuptools]
35
+ package-dir = {"" = "src"}
36
+
37
+ [tool.setuptools.packages.find]
38
+ where = ["src"]
39
+ include = ["cfskit*"]
cfskit-0.1.0/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,72 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2026/5/8 17:32
3
+ # @Author : CFuShn
4
+ # @Comments:
5
+ # @Software: PyCharm
6
+
7
+ from __future__ import annotations
8
+
9
+ from typing import Any
10
+
11
+ from cfskit.t02_log_util import setup_logger
12
+ from cfskit.t03_tb_util import setup_tensorboard, TensorBoardWritter
13
+
14
+
15
+ def _package_version() -> str:
16
+ try:
17
+ from importlib.metadata import PackageNotFoundError, version
18
+
19
+ try:
20
+ return version("cfskit")
21
+ except PackageNotFoundError:
22
+ # Editable install / source tree without metadata
23
+ return "0.1.0"
24
+ except Exception:
25
+ return "0.1.0"
26
+
27
+
28
+ __version__ = _package_version()
29
+
30
+ __all__ = [
31
+ "__version__",
32
+ # log
33
+ "setup_logger",
34
+ # tensorboard
35
+ "setup_tensorboard",
36
+ "TensorBoardWritter",
37
+ # ipc
38
+ "t04_ipc_util",
39
+ ]
40
+
41
+ _EXPORTS: dict[str, tuple[str, str]] = {
42
+ # log
43
+ "setup_logger": ("cfskit.t02_log_util", "setup_logger"),
44
+ # tensorboard
45
+ "setup_tensorboard": ("cfskit.t03_tb_util", "setup_tensorboard"),
46
+ "TensorBoardLogger": ("cfskit.t03_tb_util", "TensorBoardLogger"),
47
+ "NullTensorBoardLogger": ("cfskit.t03_tb_util", "NullTensorBoardLogger"),
48
+ # ipc
49
+ "register_signal_handler": ("cfskit.t04_ipc_util", "register_signal_handler"),
50
+ "get_s1": ("cfskit.t04_ipc_util", "get_s1"),
51
+ "get_s2": ("cfskit.t04_ipc_util", "get_s2"),
52
+ "switch_s1": ("cfskit.t04_ipc_util", "switch_s1"),
53
+ "switch_s2": ("cfskit.t04_ipc_util", "switch_s2"),
54
+ }
55
+
56
+
57
+ def __getattr__(name: str) -> Any:
58
+ """
59
+ 顶层 API 惰性导入,避免 `import cfskit` 因为可选依赖缺失而失败。
60
+ """
61
+ if name not in _EXPORTS:
62
+ raise AttributeError(f"module 'cfskit' has no attribute '{name}'")
63
+
64
+ module_name, attr_name = _EXPORTS[name]
65
+ from importlib import import_module
66
+
67
+ mod = import_module(module_name)
68
+ return getattr(mod, attr_name)
69
+
70
+
71
+ def __dir__() -> list[str]:
72
+ return sorted(set(globals().keys()) | set(_EXPORTS.keys()))
@@ -0,0 +1,138 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2026/5/8 16:13
3
+ # @Author : CFuShn
4
+ # @Comments:
5
+ # @Software: PyCharm
6
+
7
+ from __future__ import annotations
8
+
9
+ import sys
10
+ from datetime import datetime
11
+ from pathlib import Path
12
+ from typing import Optional, TYPE_CHECKING
13
+
14
+ from loguru import logger
15
+
16
+ if TYPE_CHECKING: # pragma: no cover
17
+ from loguru import Logger
18
+
19
+ _LOGGER_INITIALIZED = False
20
+ _LOG_PATH = None
21
+
22
+
23
+ class TqdmCompatibleSink:
24
+ """
25
+ 用 tqdm.write() 替代普通 print/stderr 写入,避免打乱 tqdm 进度条。
26
+ """
27
+
28
+ def __init__(self):
29
+ try:
30
+ from tqdm import tqdm
31
+ self._tqdm = tqdm
32
+ except Exception:
33
+ self._tqdm = None
34
+
35
+ def write(self, message: str):
36
+ message = message.rstrip("\n")
37
+ if not message:
38
+ return
39
+
40
+ if self._tqdm is not None:
41
+ self._tqdm.write(message)
42
+ else:
43
+ print(message, file=sys.stderr)
44
+
45
+ def flush(self):
46
+ pass
47
+
48
+
49
+ def setup_logger(
50
+ log_dir: str | Path = "logs",
51
+ name: Optional[str] = None,
52
+ level: str = "INFO",
53
+ console: bool = True,
54
+ file: bool = True,
55
+ rotation: str = "50 MB",
56
+ retention: str = "30 days",
57
+ enqueue: bool = False,
58
+ allow_reconfigure: bool = False,
59
+ ) -> tuple["Logger", Path]:
60
+ """
61
+ 项目统一 logger 初始化。
62
+
63
+ 默认行为:
64
+ - 只允许初始化一次
65
+ - 控制台彩色输出
66
+ - txt 文件纯文本输出
67
+ - INFO 级别
68
+ - tqdm 兼容
69
+ - logger.bind(console=False).info(...) 可仅写文件
70
+
71
+ 如果确实想重新配置,显式传 allow_reconfigure=True。
72
+ """
73
+
74
+ global _LOGGER_INITIALIZED, _LOG_PATH
75
+
76
+ if _LOGGER_INITIALIZED and not allow_reconfigure:
77
+ raise RuntimeError(
78
+ f"setup_logger() has already been called once. "
79
+ f"Current log file: {_LOG_PATH}. "
80
+ f"If you really want to reconfigure logger, call "
81
+ f"setup_logger(..., allow_reconfigure=True)."
82
+ )
83
+
84
+ logger.remove()
85
+
86
+ log_dir = Path(log_dir)
87
+ log_dir.mkdir(parents=True, exist_ok=True)
88
+
89
+ if name is None:
90
+ name = datetime.now().strftime("%Y%m%d_%H%M%S")
91
+
92
+ log_path = log_dir / f"{name}.txt"
93
+
94
+ console_fmt = (
95
+ "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
96
+ "<level>{level:<8}</level> | "
97
+ "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> | "
98
+ "<level>{message}</level>"
99
+ )
100
+
101
+ file_fmt = (
102
+ "{time:YYYY-MM-DD HH:mm:ss.SSS} | "
103
+ "{level:<8} | "
104
+ "{name}:{function}:{line} | "
105
+ "{message}"
106
+ )
107
+
108
+ def console_filter(record):
109
+ return record["extra"].get("console", True)
110
+
111
+ if console:
112
+ logger.add(
113
+ TqdmCompatibleSink(),
114
+ level=level,
115
+ format=console_fmt,
116
+ filter=console_filter,
117
+ colorize=True,
118
+ enqueue=enqueue,
119
+ )
120
+
121
+ if file:
122
+ logger.add(
123
+ log_path,
124
+ level=level,
125
+ format=file_fmt,
126
+ encoding="utf-8",
127
+ rotation=rotation,
128
+ retention=retention,
129
+ colorize=False,
130
+ enqueue=enqueue,
131
+ )
132
+
133
+ _LOGGER_INITIALIZED = True
134
+ _LOG_PATH = log_path
135
+
136
+ logger.info(f"Log file: {log_path}")
137
+
138
+ return logger, log_path
@@ -0,0 +1,417 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2026/5/8 16:13
3
+ # @Author : CFuShn
4
+ # @Comments:
5
+ # @Software: PyCharm
6
+
7
+ from __future__ import annotations
8
+
9
+ import json
10
+ import socket
11
+ from datetime import datetime
12
+ from pathlib import Path
13
+ from typing import Any, Optional, Mapping, Union, TYPE_CHECKING
14
+
15
+ if TYPE_CHECKING: # pragma: no cover
16
+ import torch
17
+ from torch.utils.tensorboard import SummaryWriter
18
+
19
+ _TB_INITIALIZED = False
20
+
21
+
22
+ def _is_main_process(accelerator=None) -> bool:
23
+ if accelerator is None:
24
+ return True
25
+ return bool(accelerator.is_main_process)
26
+
27
+
28
+ def _to_float(value: Any) -> float:
29
+ """
30
+ 尽量把 torch scalar / numpy scalar / python number 转成 float。
31
+ """
32
+ try:
33
+ import torch # type: ignore
34
+
35
+ if isinstance(value, torch.Tensor):
36
+ value = value.detach()
37
+ if value.numel() != 1:
38
+ raise ValueError(f"Expected scalar tensor, got shape={tuple(value.shape)}")
39
+ return float(value.item())
40
+ except Exception:
41
+ pass
42
+ return float(value)
43
+
44
+
45
+ def _to_device_cpu(x: Any):
46
+ """
47
+ TensorBoard 写图片/直方图时,尽量 detach + cpu,避免持有计算图或 GPU 显存。
48
+ """
49
+ try:
50
+ import torch # type: ignore
51
+
52
+ if isinstance(x, torch.Tensor):
53
+ return x.detach().cpu()
54
+ except Exception:
55
+ pass
56
+ return x
57
+
58
+
59
+ class NullTensorBoardWritter:
60
+ """
61
+ 非主进程使用的空操作 logger。
62
+ 这样训练代码可以无脑写 tb.scalar(...),不用到处 if rank0。
63
+ """
64
+
65
+ enabled = False
66
+ log_dir = None
67
+ writer = None
68
+
69
+ def __getattr__(self, name):
70
+ def noop(*args, **kwargs):
71
+ return None
72
+
73
+ return noop
74
+
75
+ def close(self):
76
+ return None
77
+
78
+ def flush(self):
79
+ return None
80
+
81
+
82
+ class TensorBoardWritter:
83
+ """
84
+ 对 SummaryWriter 的轻量封装。
85
+
86
+ 常用方法:
87
+ - scalar
88
+ - scalars
89
+ - metrics
90
+ - image
91
+ - images
92
+ - histogram
93
+ - histograms
94
+ - text
95
+ - config
96
+ - graph
97
+ - flush
98
+ - close
99
+ """
100
+
101
+ enabled = True
102
+
103
+ def __init__(
104
+ self,
105
+ log_dir: str | Path,
106
+ *,
107
+ max_queue: int = 10,
108
+ flush_secs: int = 120,
109
+ comment: str = "",
110
+ filename_suffix: str = "",
111
+ ):
112
+ try:
113
+ from torch.utils.tensorboard import SummaryWriter # type: ignore
114
+ except Exception as e: # pragma: no cover
115
+ raise ImportError(
116
+ "TensorBoardLogger 需要 PyTorch。请执行:pip install \"cfskit[torch]\"(或 cfskit[tb])。"
117
+ ) from e
118
+
119
+ self.log_dir = Path(log_dir)
120
+ self.log_dir.mkdir(parents=True, exist_ok=True)
121
+
122
+ self.writer: "SummaryWriter" = SummaryWriter(
123
+ log_dir=str(self.log_dir),
124
+ comment=comment,
125
+ max_queue=max_queue,
126
+ flush_secs=flush_secs,
127
+ filename_suffix=filename_suffix,
128
+ )
129
+
130
+ def scalar(self, tag: str, value: Any, step: int):
131
+ self.writer.add_scalar(tag, _to_float(value), global_step=step)
132
+
133
+ def scalars(self, main_tag: str, tag_scalar_dict: Mapping[str, Any], step: int):
134
+ values = {k: _to_float(v) for k, v in tag_scalar_dict.items()}
135
+ self.writer.add_scalars(main_tag, values, global_step=step)
136
+
137
+ def metrics(
138
+ self,
139
+ metrics: Mapping[str, Any],
140
+ step: int,
141
+ *,
142
+ prefix: Optional[str] = None,
143
+ ):
144
+ """
145
+ 批量写 scalar。
146
+
147
+ 示例:
148
+ tb.metrics({"loss": loss, "lr": lr}, step, prefix="Train")
149
+ 写成:
150
+ Train/loss
151
+ Train/lr
152
+ """
153
+ for k, v in metrics.items():
154
+ tag = f"{prefix}/{k}" if prefix else k
155
+ self.scalar(tag, v, step)
156
+
157
+ def image(
158
+ self,
159
+ tag: str,
160
+ img_tensor: "torch.Tensor",
161
+ step: int,
162
+ *,
163
+ dataformats: str = "CHW",
164
+ ):
165
+ """
166
+ 写单张图片。
167
+
168
+ 常见格式:
169
+ - CHW: [C, H, W]
170
+ - HWC: [H, W, C]
171
+ """
172
+ img_tensor = _to_device_cpu(img_tensor)
173
+ self.writer.add_image(
174
+ tag,
175
+ img_tensor,
176
+ global_step=step,
177
+ dataformats=dataformats,
178
+ )
179
+
180
+ def images(
181
+ self,
182
+ tag: str,
183
+ img_tensor: "torch.Tensor",
184
+ step: int,
185
+ *,
186
+ dataformats: str = "NCHW",
187
+ ):
188
+ """
189
+ 写一组图片。
190
+
191
+ 常见格式:
192
+ - NCHW: [N, C, H, W]
193
+ - NHWC: [N, H, W, C]
194
+ """
195
+ img_tensor = _to_device_cpu(img_tensor)
196
+ self.writer.add_images(
197
+ tag,
198
+ img_tensor,
199
+ global_step=step,
200
+ dataformats=dataformats,
201
+ )
202
+
203
+ def image_grid(
204
+ self,
205
+ tag: str,
206
+ img_tensor: "torch.Tensor",
207
+ step: int,
208
+ *,
209
+ nrow: int = 4,
210
+ normalize: bool = True,
211
+ value_range: Optional[tuple[float, float]] = None,
212
+ ):
213
+ """
214
+ 把 [B, C, H, W] 做成 grid 后写入。
215
+ SR / ContSR 任务非常常用。
216
+ """
217
+ from torchvision.utils import make_grid
218
+
219
+ img_tensor = _to_device_cpu(img_tensor)
220
+
221
+ grid = make_grid(
222
+ img_tensor,
223
+ nrow=nrow,
224
+ normalize=normalize,
225
+ value_range=value_range,
226
+ )
227
+ self.image(tag, grid, step, dataformats="CHW")
228
+
229
+ def sr_triplet(
230
+ self,
231
+ tag: str,
232
+ lr: "torch.Tensor",
233
+ sr: "torch.Tensor",
234
+ hr: "torch.Tensor",
235
+ step: int,
236
+ *,
237
+ nrow: int = 4,
238
+ normalize: bool = True,
239
+ with_error: bool = True,
240
+ ):
241
+ """
242
+ 超分任务常用:一次性记录 LR / SR / HR / Error Map。
243
+
244
+ 要求:
245
+ - lr / sr / hr 尽量都是 [B, C, H, W]
246
+ - 如果 LR 尺寸小于 SR/HR,建议你在外面先 interpolate 到相同尺寸
247
+ """
248
+ self.image_grid(f"{tag}/lr", lr, step, nrow=nrow, normalize=normalize)
249
+ self.image_grid(f"{tag}/sr", sr, step, nrow=nrow, normalize=normalize)
250
+ self.image_grid(f"{tag}/hr", hr, step, nrow=nrow, normalize=normalize)
251
+
252
+ if with_error:
253
+ err = (sr.detach() - hr.detach()).abs()
254
+ self.image_grid(f"{tag}/error", err, step, nrow=nrow, normalize=True)
255
+
256
+ def histogram(self, tag: str, values: "torch.Tensor", step: int, *, bins: str = "tensorflow"):
257
+ values = _to_device_cpu(values)
258
+ self.writer.add_histogram(tag, values, global_step=step, bins=bins)
259
+
260
+ def histograms(
261
+ self,
262
+ model: "torch.nn.Module",
263
+ step: int,
264
+ *,
265
+ log_weights: bool = True,
266
+ log_grads: bool = True,
267
+ prefix: str = "",
268
+ ):
269
+ """
270
+ 记录模型权重和梯度分布。
271
+ 注意:不要每个 step 都写,建议几百/几千 step 写一次。
272
+ """
273
+ prefix = prefix.strip("/")
274
+ for name, param in model.named_parameters():
275
+ safe_name = name.replace(".", "/")
276
+ if prefix:
277
+ safe_name = f"{prefix}/{safe_name}"
278
+
279
+ if log_weights:
280
+ self.histogram(f"Weights/{safe_name}", param, step)
281
+
282
+ if log_grads and param.grad is not None:
283
+ self.histogram(f"Gradients/{safe_name}", param.grad, step)
284
+
285
+ def grad_norms(
286
+ self,
287
+ model: "torch.nn.Module",
288
+ step: int,
289
+ *,
290
+ prefix: str = "GradNorm",
291
+ norm_type: float = 2.0,
292
+ ):
293
+ """
294
+ 记录每个参数的梯度范数。
295
+ 比 histogram 更轻量。
296
+ """
297
+ for name, param in model.named_parameters():
298
+ if param.grad is None:
299
+ continue
300
+ tag = f"{prefix}/{name.replace('.', '/')}"
301
+ value = param.grad.detach().norm(norm_type)
302
+ self.scalar(tag, value, step)
303
+
304
+ def text(self, tag: str, text_string: str, step: int = 0):
305
+ self.writer.add_text(tag, text_string, global_step=step)
306
+
307
+ def config(self, config: Any, step: int = 0, tag: str = "Config"):
308
+ """
309
+ 写入配置。
310
+ 支持 dict / list / str / 其他可 json 化对象。
311
+ """
312
+ if isinstance(config, str):
313
+ text = config
314
+ else:
315
+ try:
316
+ text = json.dumps(config, indent=2, ensure_ascii=False)
317
+ text = f"```json\n{text}\n```"
318
+ except TypeError:
319
+ text = f"```text\n{str(config)}\n```"
320
+
321
+ self.text(tag, text, step=step)
322
+
323
+ def graph(self, model: "torch.nn.Module", input_to_model: Any):
324
+ """
325
+ 记录模型 graph。
326
+ 对动态 forward / 多输入模型可能失败,这是 TensorBoard graph 本身限制。
327
+ """
328
+ try:
329
+ import torch # type: ignore
330
+
331
+ if isinstance(input_to_model, torch.Tensor):
332
+ input_to_model = input_to_model.detach()
333
+ elif isinstance(input_to_model, tuple):
334
+ input_to_model = tuple(
335
+ x.detach() if isinstance(x, torch.Tensor) else x
336
+ for x in input_to_model
337
+ )
338
+ except Exception:
339
+ pass
340
+ self.writer.add_graph(model, input_to_model)
341
+
342
+ def hparams(
343
+ self,
344
+ hparam_dict: Mapping[str, Any],
345
+ metric_dict: Mapping[str, Any],
346
+ ):
347
+ """
348
+ 记录超参数。
349
+ 注意:SummaryWriter.add_hparams 会生成额外子目录,TensorBoard 里显示方式和普通 scalar 不完全一样。
350
+ """
351
+ clean_metrics = {k: _to_float(v) for k, v in metric_dict.items()}
352
+ self.writer.add_hparams(dict(hparam_dict), clean_metrics)
353
+
354
+ def flush(self):
355
+ self.writer.flush()
356
+
357
+ def close(self):
358
+ self.writer.flush()
359
+ self.writer.close()
360
+
361
+
362
+ def setup_tensorboard(
363
+ log_dir: str | Path = "runs",
364
+ name: Optional[str] = None,
365
+ *,
366
+ project: Optional[str] = None,
367
+ run_id: Optional[str] = None,
368
+ accelerator=None,
369
+ max_queue: int = 10,
370
+ flush_secs: int = 120,
371
+ comment: str = "",
372
+ filename_suffix: str = "",
373
+ allow_reconfigure: bool = False,
374
+ ) -> Union["TensorBoardWritter", "NullTensorBoardWritter"]:
375
+ global _TB_INITIALIZED
376
+
377
+ if _TB_INITIALIZED and not allow_reconfigure:
378
+ raise RuntimeError(
379
+ "setup_tensorboard() has already been called once. "
380
+ "If you really want to reconfigure it, call "
381
+ "setup_tensorboard(..., allow_reconfigure=True)."
382
+ )
383
+
384
+ is_main = _is_main_process(accelerator)
385
+
386
+ base_dir = Path(log_dir)
387
+
388
+ if run_id is not None:
389
+ run_name = run_id
390
+ else:
391
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
392
+ hostname = socket.gethostname()
393
+
394
+ if name is None:
395
+ run_name = f"{timestamp}_{hostname}"
396
+ else:
397
+ run_name = f"{timestamp}_{hostname}_{name}"
398
+
399
+ if project is not None:
400
+ tb_log_dir = base_dir / project / run_name
401
+ else:
402
+ tb_log_dir = base_dir / run_name
403
+
404
+ if is_main:
405
+ tb = TensorBoardWritter(
406
+ tb_log_dir,
407
+ max_queue=max_queue,
408
+ flush_secs=flush_secs,
409
+ comment=comment,
410
+ filename_suffix=filename_suffix,
411
+ )
412
+ else:
413
+ tb = NullTensorBoardWritter()
414
+
415
+ _TB_INITIALIZED = True
416
+
417
+ return tb
@@ -0,0 +1,83 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2025/2/12 13:50
3
+ # @Author : CFuShn
4
+ # @Comments:
5
+ # @Software: PyCharm
6
+
7
+ import os
8
+ import signal
9
+ import threading
10
+
11
+ # 两个自定义信号量, 供用户处理事件, 通过SIGUSR1 和 SIGUSR2 信号操作其值
12
+ __s1 = 0 # 对应SIGUSR1 = 10
13
+ __s2 = 0 # 对应SIGUSR2 = 12
14
+ # 全局锁, 保护 __s1 / __s2 的读写
15
+ __sig_lock = threading.Lock()
16
+
17
+
18
+ def _signal_handler(signum, frame):
19
+ msg = None
20
+ global __s1, __s2
21
+ with __sig_lock:
22
+ if signum == signal.SIGUSR1:
23
+ __s1 = 1 - __s1 # 在0,1之间switch
24
+ msg = f"\n================= Received signal {signum}, set s1 = {__s1} ==================\n"
25
+ elif signum == signal.SIGUSR2:
26
+ __s2 = 1 - __s2 # 在0,1之间switch
27
+ msg = f"\n================= Received signal {signum}, set s2 = {__s2} ==================\n"
28
+ msg and print(msg)
29
+
30
+
31
+ def register_signal_handler():
32
+ """注册用户自定义signal_handler,在接收信号时触发触发对应信号量的自定义操作"""
33
+ signal.signal(signal.SIGUSR1, _signal_handler)
34
+ signal.signal(signal.SIGUSR2, _signal_handler)
35
+ print(f"================= signal handler for pid: {os.getpid()} registered. =================")
36
+
37
+
38
+ def get_s1():
39
+ """对应信号:SIGUSR1 = 10"""
40
+ with __sig_lock:
41
+ return __s1
42
+
43
+
44
+ def get_s2():
45
+ """对应信号:SIGUSR2 = 12"""
46
+ with __sig_lock:
47
+ return __s2
48
+
49
+
50
+ def switch_s1(num):
51
+ """对应信号:SIGUSR1 = 10"""
52
+ assert num in (0, 1)
53
+ global __s1
54
+ with __sig_lock: # 在 switch 处加安全锁
55
+ if __s1 != num:
56
+ __s1 = num
57
+ print(f"=================== switch s1 to {__s1} ====================")
58
+
59
+
60
+ def switch_s2(num):
61
+ """对应信号:SIGUSR2 = 12"""
62
+ assert num in (0, 1)
63
+ global __s2
64
+ with __sig_lock: # 在 switch 处加安全锁
65
+ if __s2 != num:
66
+ __s2 = num
67
+ print(f"=================== switch s2 to {__s2} ====================")
68
+
69
+
70
+ if __name__ == "__main__":
71
+
72
+ import sys
73
+ import time
74
+
75
+ # 必须先注册, 信号操作才会生效
76
+ register_signal_handler()
77
+
78
+ for i in range(50):
79
+ get_s2() > 0 and sys.exit("SIGUSR2 is active! This process exit!")
80
+ for j in range(10):
81
+ time.sleep(1)
82
+ get_s1() and print("SIGUSR1 is active!")
83
+ print(i, j)
@@ -0,0 +1,5 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2026/5/8 18:43
3
+ # @Author : CFuShn
4
+ # @Comments:
5
+ # @Software: PyCharm
@@ -0,0 +1,65 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2025/5/23 16:08
3
+ # @Author : CFuShn
4
+ # @Comments: 全局上下文文件,用于加载配置信息,统管上下文
5
+ # @Software: PyCharm
6
+
7
+ try:
8
+ import tomllib # Python 3.11+
9
+ except ModuleNotFoundError:
10
+ import tomli as tomllib # For Python <3.11
11
+ from dataclasses import dataclass
12
+ from typing import Optional
13
+
14
+ from accelerate.utils import set_seed
15
+
16
+
17
+ @dataclass
18
+ class CFuShnContext:
19
+ # sys & env
20
+ seed = 3407 # 3407 is all you need
21
+
22
+ # 系统,环境
23
+ project_root = ""
24
+ output_root = ""
25
+ ckpt_version = ""
26
+
27
+ # 上下文传递
28
+ # ...
29
+
30
+ def __str__(self) -> str:
31
+ """一行一行打印参数"""
32
+ pretty = ""
33
+ # for key, value in self.__dict__.items():
34
+ # pretty += (f"{key}: {value}\n" if key != "_instance" else "")
35
+ # 上面的方法无法保证属性是有序输出的,
36
+ # 使用__annotations__获取成员变量就是有序的, 但是在定义时必须指定其类型
37
+ # eg: version = xxx 就不在__annotations__中, 必须指定类型: version:str = xxx
38
+ for key in self.__annotations__.keys():
39
+ pretty += (f"{key}: {getattr(self, key)}\n" if key != "_instance" else "")
40
+ return pretty
41
+
42
+ # 单例实例
43
+ _instance: Optional["CFuShnContext"] = None
44
+ _initialized = False # 初始化保护
45
+
46
+ def __new__(cls):
47
+ if cls._instance is None:
48
+ cls._instance = super().__new__(cls)
49
+ return cls._instance
50
+
51
+ def __init__(self):
52
+ """
53
+ 整个项目基础上下文的初始化
54
+ """
55
+ if self._initialized:
56
+ return # 防止重复初始化
57
+ self._initialized = True
58
+
59
+ # 设置所有相关的的随机数种子
60
+ # use accelerate lib # set the seed in `random`, `numpy`, `torch`
61
+ set_seed(self.seed)
62
+
63
+
64
+ # 实例化并作为全局配置对象
65
+ ctx = CFuShnContext()
@@ -0,0 +1,5 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2026/5/8 17:30
3
+ # @Author : CFuShn
4
+ # @Comments:
5
+ # @Software: PyCharm
@@ -0,0 +1,61 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2025/6/3 21:11
3
+ # @Author : CFuShn
4
+ # @Comments: 单张人脸图像
5
+ # @Software: PyCharm
6
+
7
+ from typing import Optional
8
+
9
+ import cv2
10
+
11
+
12
+ class CcImage:
13
+
14
+ def __init__(self, _id, path=None, origin=None):
15
+ # 基本信息
16
+ self.id: str = _id # 必传 (格式: 数据集id_{path.basename})
17
+ self.path: Optional[str | None] = path
18
+ self.origin: Optional[str | None] = origin # 溯源
19
+
20
+ # 文本prompt
21
+ self.prompt: str = "" # prompt 或 image caption
22
+
23
+ def img(self, mode="BGR"):
24
+ # img = cv2.imread(self.path, cv2.IMREAD_UNCHANGED) # IMREAD_UNCHANGED表示保留alpha通道
25
+ bgr_img = cv2.imread(self.path) # 默认是IMREAD_COLOR模式,只读取彩色, 且默认返回BGR顺序
26
+ if bgr_img is None:
27
+ return None
28
+ img = _ensure_bgr2rgb(bgr_img) if mode == "RGB" else _ensure_bgr(bgr_img)
29
+ return img
30
+
31
+
32
+ # 将 BGR图像转换为RGB, 同时确保是彩色
33
+ def _ensure_bgr2rgb(img):
34
+ if img is None:
35
+ return None
36
+ if len(img.shape) == 2: # 灰度图
37
+ return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
38
+ elif img.shape[2] == 1: # 单通道
39
+ return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
40
+ elif img.shape[2] == 3: # BGR->RGB
41
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
42
+ return np.clip(img, 0, 255).astype('uint8') # 新增保证
43
+ else:
44
+ print(f"[WARNING] 非标准图像通道数: {img.shape}")
45
+ return img
46
+
47
+
48
+ def _ensure_bgr(img):
49
+ if img is None:
50
+ return None
51
+ if len(img.shape) == 2: # 灰度图
52
+ return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
53
+ elif img.shape[2] == 1: # 单通道
54
+ return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
55
+ elif img.shape[2] == 3: # 已是BGR
56
+ return img
57
+ elif img.shape[2] == 4: # BGRA -> BGR
58
+ return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
59
+ else:
60
+ print(f"[WARNING] 非标准图像通道数: {img.shape}")
61
+ return img[:, :, :3] # 默认保前三通道
@@ -0,0 +1,110 @@
1
+ # Differentiable Augmentation for Data-Efficient GAN Training
2
+ # Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song Han
3
+ # https://arxiv.org/pdf/2006.10738
4
+
5
+ import torch
6
+ import torch.nn.functional as F
7
+
8
+
9
+ def DiffAugment(x, policy='', channels_first=True):
10
+ """
11
+ 在 GAN 训练过程中进行数据增强,适用于数据量较小的情况,以防止模式崩溃(mode collapse)。
12
+ 增强操作是可微分的,确保反向传播不会被打断。
13
+ 支持不同的增强策略(policy),如亮度、对比度、平移、裁剪等。
14
+ """
15
+ if policy:
16
+ if not channels_first:
17
+ x = x.permute(0, 3, 1, 2)
18
+ for p in policy.split(','):
19
+ for f in AUGMENT_FNS[p]:
20
+ x = f(x)
21
+ if not channels_first:
22
+ x = x.permute(0, 2, 3, 1)
23
+ x = x.contiguous()
24
+ return x
25
+
26
+
27
+ def rand_brightness(x):
28
+ """
29
+ 亮度增强
30
+ 给图像整体加上一个随机亮度偏移量 [-0.5, 0.5]
31
+ (由于 rand_brightness 直接对 x 进行加法运算,因此是可微分的)
32
+ """
33
+ x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)
34
+ return x
35
+
36
+
37
+ def rand_saturation(x):
38
+ """
39
+ 饱和度增强
40
+ 计算 x 在 通道维度 的均值(即去除颜色信息,转灰度)。
41
+ 按一个随机因子 缩放颜色饱和度。
42
+ """
43
+ x_mean = x.mean(dim=1, keepdim=True)
44
+ x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean
45
+ return x
46
+
47
+
48
+ def rand_contrast(x):
49
+ """
50
+ 对比度增强
51
+ 计算整个图像 x 的均值;
52
+ 让 x 按随机因子缩放,使得对比度变化。
53
+ """
54
+ x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
55
+ x = (x - x_mean) * (
56
+ torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean
57
+ return x
58
+
59
+
60
+ def rand_translation(x, ratio=0.125):
61
+ """
62
+ 随机平移
63
+ 使其在一定范围内进行随机像素偏移,防止模型过拟合于固定位置的特征。
64
+ 适用于 GAN 训练,使生成器学习到更丰富的特征,而不是依赖固定的位置关系。
65
+ 通过 meshgrid + 索引变换,避免插值影响可微性
66
+ """
67
+ # ratio=0.125 表示平移范围是图片尺寸的 12.5%。
68
+ shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
69
+ translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
70
+ translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
71
+ grid_batch, grid_x, grid_y = torch.meshgrid(
72
+ torch.arange(x.size(0), dtype=torch.long, device=x.device),
73
+ torch.arange(x.size(2), dtype=torch.long, device=x.device),
74
+ torch.arange(x.size(3), dtype=torch.long, device=x.device),
75
+ )
76
+ grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
77
+ grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
78
+ x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
79
+ x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
80
+ return x
81
+
82
+
83
+ def rand_cutout(x, ratio=0.5):
84
+ """
85
+ 随机遮挡一部分图像区域
86
+ 让模型学会忽略局部信息,通过 mask 操作高效实现裁剪区域零填充
87
+ """
88
+ cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
89
+ offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1],
90
+ device=x.device)
91
+ offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1],
92
+ device=x.device)
93
+ grid_batch, grid_x, grid_y = torch.meshgrid(
94
+ torch.arange(x.size(0), dtype=torch.long, device=x.device),
95
+ torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
96
+ torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
97
+ )
98
+ grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)
99
+ grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)
100
+ mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
101
+ mask[grid_batch, grid_x, grid_y] = 0
102
+ x = x * mask.unsqueeze(1)
103
+ return x
104
+
105
+
106
+ AUGMENT_FNS = {
107
+ 'color': [rand_brightness, rand_saturation, rand_contrast],
108
+ 'translation': [rand_translation],
109
+ 'cutout': [rand_cutout],
110
+ }
@@ -0,0 +1,5 @@
1
+ # -*- coding: utf-8 -*-
2
+ # @Time : 2025/10/30 13:55
3
+ # @Author : CFuShn
4
+ # @Comments:
5
+ # @Software: PyCharm
@@ -0,0 +1,44 @@
1
+ Metadata-Version: 2.4
2
+ Name: cfskit
3
+ Version: 0.1.0
4
+ Summary: Common research utilities for our NUIST-GenAI-Lab and CCGM tasks
5
+ Author-email: CFuShn <1354809038@qq.com>
6
+ License-Expression: MIT
7
+ Requires-Python: >=3.8
8
+ Description-Content-Type: text/markdown
9
+ License-File: LICENSE
10
+ Requires-Dist: loguru<1,>=0.6.0
11
+ Requires-Dist: accelerate<1,>=0.23.0; python_version < "3.10"
12
+ Requires-Dist: accelerate<2,>=1.0.0; python_version >= "3.10"
13
+ Requires-Dist: tensorboard<2.20,>=2.12.0; python_version < "3.9"
14
+ Requires-Dist: tensorboard<3,>=2.20.0; python_version >= "3.9"
15
+ Provides-Extra: torch
16
+ Requires-Dist: torch<3,>=2.0.0; extra == "torch"
17
+ Provides-Extra: tb
18
+ Requires-Dist: torch<3,>=2.0.0; extra == "tb"
19
+ Provides-Extra: tb-vision
20
+ Requires-Dist: torchvision<1,>=0.15.0; extra == "tb-vision"
21
+ Dynamic: license-file
22
+
23
+ # cfskit
24
+ Common research utilities for our NUIST-GenAI-Lab and CCGM tasks
25
+
26
+
27
+
28
+ ## Features
29
+
30
+ - 整合常用的科研库, 减少工具学习成本
31
+ - 尽量规范
32
+ - 尽量简洁
33
+ - 约定大于配置
34
+ - 最小化上手难度
35
+
36
+
37
+
38
+ ## 整合功能
39
+
40
+ - accelerate封装
41
+ - tensorboard封装
42
+ - loguru封装
43
+ - ipc 信号中断工具
44
+ - 常用训练,评估,测试工具
@@ -0,0 +1,18 @@
1
+ LICENSE
2
+ README.md
3
+ pyproject.toml
4
+ src/cfskit/__init__.py
5
+ src/cfskit/t02_log_util.py
6
+ src/cfskit/t03_tb_util.py
7
+ src/cfskit/t04_ipc_util.py
8
+ src/cfskit/version.py
9
+ src/cfskit.egg-info/PKG-INFO
10
+ src/cfskit.egg-info/SOURCES.txt
11
+ src/cfskit.egg-info/dependency_links.txt
12
+ src/cfskit.egg-info/requires.txt
13
+ src/cfskit.egg-info/top_level.txt
14
+ src/cfskit/todos/__init__.py
15
+ src/cfskit/todos/t01_context.py
16
+ src/cfskit/todos/t31_omni_dataset.py
17
+ src/cfskit/todos/t32_cc_image.py
18
+ src/cfskit/todos/t90_DiffAugment_pytorch.py
@@ -0,0 +1,22 @@
1
+ loguru<1,>=0.6.0
2
+
3
+ [:python_version < "3.10"]
4
+ accelerate<1,>=0.23.0
5
+
6
+ [:python_version < "3.9"]
7
+ tensorboard<2.20,>=2.12.0
8
+
9
+ [:python_version >= "3.10"]
10
+ accelerate<2,>=1.0.0
11
+
12
+ [:python_version >= "3.9"]
13
+ tensorboard<3,>=2.20.0
14
+
15
+ [tb]
16
+ torch<3,>=2.0.0
17
+
18
+ [tb-vision]
19
+ torchvision<1,>=0.15.0
20
+
21
+ [torch]
22
+ torch<3,>=2.0.0
@@ -0,0 +1 @@
1
+ cfskit