feat(cluster): 加 --train-only,自监督训练与聚类可分两步跑
engine/simclr.py 与 engine/cluster.py 本来就互不引用,唯一接口是 simclr_backbone.pth;耦合只在 CLI——method=simclr 会训练完直接聚类。 此前"只聚类"可以用 simclr.checkpoint 做到,"只训练"没有办法。 - --train-only:训练存完权重即退出,不抽特征也不聚类,并打印下一步的聚类命令 - 与 method=pretrained / simclr.checkpoint / --sweep-k 组合时立刻报错说明原因, 而不是跑完才发现没有产出 - 典型用法:GPU 机器上 --train-only 跑几十小时,之后拿权重反复调 k,每次约 1 分钟 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
5d395b7c25
commit
0f5bb4a5c0
@@ -389,6 +389,11 @@ def main() -> int:
|
||||
parser.add_argument(
|
||||
"--sweep-k", default=None, metavar="A-B", help="只扫簇数(如 2-8)打印指标表,不产出分组"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--train-only",
|
||||
action="store_true",
|
||||
help="只做 SimCLR 自监督训练,存完权重就退出,不抽特征也不聚类",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
cfg = load_config(args.config, args.set)
|
||||
@@ -409,18 +414,46 @@ def main() -> int:
|
||||
logger.error("%s", exc)
|
||||
return 1
|
||||
|
||||
method = str(cfg.get("method", "pretrained"))
|
||||
if args.train_only:
|
||||
# 这三种组合下 --train-only 要么无事可做、要么与另一个"提前结束"的开关打架,
|
||||
# 与其跑完才发现什么都没产出,不如立刻说清楚
|
||||
if method != "simclr":
|
||||
logger.error("--train-only 只对 method=simclr 有意义,当前 method=%s(没有要训练的东西)", method)
|
||||
return 1
|
||||
if cfg.get("simclr", {}).get("checkpoint"):
|
||||
logger.error(
|
||||
"--train-only 与 simclr.checkpoint 冲突:checkpoint 会跳过训练,"
|
||||
"加上 --train-only 就什么都不会做。想接着上次练请改用 simclr.resume。"
|
||||
)
|
||||
return 1
|
||||
if args.sweep_k:
|
||||
logger.error("--train-only 与 --sweep-k 不能同时用:前者不聚类,后者要靠聚类才能扫")
|
||||
return 1
|
||||
|
||||
dataset = build_dataset(dict(cfg["dataset"]))
|
||||
logger.info(
|
||||
"扫描到 %d 张图,来自 %d 个目录;簇数 %d,method=%s,设备 %s",
|
||||
len(dataset),
|
||||
len({dataset.source_dir(path) for path in dataset.images}),
|
||||
num_clusters,
|
||||
cfg.get("method", "pretrained"),
|
||||
method,
|
||||
describe_device(device),
|
||||
)
|
||||
|
||||
encoder = build_encoder(cfg, device, work_dir)
|
||||
|
||||
if args.train_only:
|
||||
weights = work_dir / "simclr_backbone.pth"
|
||||
logger.info("--train-only:训练完成,未做聚类。权重在 %s", weights)
|
||||
logger.info(
|
||||
"接着聚类:uv run python scripts/cluster_unlabeled.py --config %s "
|
||||
"--work-dir <新目录> --set simclr.checkpoint=%s",
|
||||
cfg.get("config_path", "<配置>"),
|
||||
weights,
|
||||
)
|
||||
return 0
|
||||
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
loader = DataLoader(
|
||||
|
||||
Reference in New Issue
Block a user