diff --git a/README.md b/README.md index 876ec37..84221f1 100644 --- a/README.md +++ b/README.md @@ -1917,6 +1917,13 @@ open-stream replay path uses version-specific Thinker internals. On Linux, compatible NVIDIA GPU, optionally install `flash-attn` and pass `--flash-attention-2`. +The separate B1 local extractor, `extract_turn_local_features.py`, is tested +with both Transformers 4.52.3 and 4.57.6. An existing 4.57.6 GPU container does +not need to be downgraded for local extraction. This does not change the replay +API's dependency assumptions. Keep all local sidecars for one experiment on +the same version; their metadata records the actual runtime version. See +[the local feature guide](docs/local_feature_baseline.md). + ## Ego4D Audio-Visual Diarization download Do not download the complete Ego4D release for this experiment. Download only diff --git a/docs/local_feature_baseline.md b/docs/local_feature_baseline.md index 5ed114a..19614d4 100644 --- a/docs/local_feature_baseline.md +++ b/docs/local_feature_baseline.md @@ -28,9 +28,21 @@ Qwen 3B 的 **2048 是编码器原生输出维度，256 只是新增 head 的可 这保证相对于 readout 时刻的输入因果性，但不保证编码器在窗口内部逐 token 因果。 实现依据：[Qwen 3B 配置](https://huggingface.co/Qwen/Qwen2.5-Omni-3B/blob/main/config.json)、 -[Transformers 4.52.3 编码器 API](https://github.com/huggingface/transformers/blob/v4.52.3/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py)。 +[Transformers 4.52.3 编码器 API](https://github.com/huggingface/transformers/blob/v4.52.3/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py)、 +[4.57.6 编码器 API](https://github.com/huggingface/transformers/blob/v4.57.6/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py)。 脚本加载 Thinker checkpoint 后释放语言解码器，仅将冻结的视觉／音频塔移到 GPU；没有 Thinker forward。 +局部提取已适配 **Transformers 4.52.3 和 4.57.6**，现有 L20 / ms-swift 容器若是 4.57.6， +无需降级，也不要为这一步重新安装整个 `requirements.txt`。旧脚本只接受精确的 4.52.3， +因此即使接口存在也会在加载模型前主动报错；更新脚本后重跑原 smoke 命令即可。 +两个版本都校验原生 token 输出形状与权重完整性，并显式使用 `use_fast=False` 固定图像预处理路径。 +未知版本仍在媒体解码前报错；`--preview-only` 不受编码器版本检查影响。 +该兼容范围只针对本局部提取脚本，不代表 Thinker KV replay 或 ROMA 的全部接口可互换。 + +输出的 `config.transformers` 记录实际运行版本，不再写死 4.52.3。 +同一局部特征实验的 train/test sessions 应统一版本、权重和提取配置；加载器会拒绝混用版本。 +接口兼容不等于不同版本的浮点特征逐位相同，已有全局 Omni 特征无需重提。 + ## 1. EasyCom：先检查框，再跑一个 4s smoke test 在已有项目依赖的 GPU 机器运行，以下路径沿用之前实验；raw 数据根按实际挂载修改。 @@ -83,6 +95,8 @@ python - <<'PY' import torch p = torch.load('features/local_smoke/Session_1.pt', map_location='cpu', weights_only=False) print('partial:', p['partial']) +print('transformers:', p['config']['transformers']) +print('box coordinates:', p['config']['box_coordinate_space']) for key in ('visual', 'audio'): x = p[key] print(key, tuple(x.shape), x.dtype, 'finite:', bool(torch.isfinite(x).all()), @@ -94,9 +108,12 @@ PY 3B 预期 V/A 各 `[20,2048]`，时间 0.2–4.0s；没有人脸的时刻允许 V 无效。 本机已验证 CPU 对齐／融合／合成数据训练推理，真实 checkpoint + GPU 提取仍需此 smoke test。 -另有不下载权重的小型随机 Qwen 编码器 API 测试。可在已有依赖的环境运行： +另有不下载权重的小型随机 Qwen 编码器 API 和本地 checkpoint 加载测试， +在 4.52.3 / 4.57.6 上均适用；checkpoint 加载测试需要 `accelerate`。 +可先检查本机编码器接口，再跑完整回归： ```bash +python -m unittest -v test_local_turn_features.TinyQwenEncoderTests python -m unittest -v test_easycom_head_objectives test_local_turn_features ``` diff --git a/extract_turn_local_features.py b/extract_turn_local_features.py index 4313caa..12ed97f 100644 --- a/extract_turn_local_features.py +++ b/extract_turn_local_features.py @@ -26,6 +26,20 @@ from local_turn_features import ( ) +TESTED_TRANSFORMERS_VERSIONS = ("4.52.3", "4.57.6") + + +def validate_local_transformers_version(version: str) -> str: + """Keep the local encoder's tested API separate from the replay KV API.""" + if version not in TESTED_TRANSFORMERS_VERSIONS: + supported = ", ".join(TESTED_TRANSFORMERS_VERSIONS) + raise RuntimeError( + f"Local encoder API is tested with transformers versions {supported}; found {version}. " + "Use a tested version or validate the encoder APIs before adding support." + ) + return version + + def model_config_digest(config: dict[str, Any]) -> str: """Ignore location fields so an explicitly identified checkpoint can move.""" def stable(value): @@ -45,12 +59,13 @@ class FrozenQwenLocalEncoder: import transformers from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration - if transformers.__version__ != "4.52.3": - raise RuntimeError("Local encoder API is tested against transformers==4.52.3; use requirements.txt") + self.transformers_version = validate_local_transformers_version(transformers.__version__) + print(f"Local encoders: transformers={self.transformers_version}", flush=True) self.device = torch.device(args.device) if self.device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError(f"CUDA requested but unavailable: {args.device}") - kwargs = {"torch_dtype": getattr(torch, args.dtype), "device_map": {"": "cpu"}, + dtype_key = "dtype" if self.transformers_version == "4.57.6" else "torch_dtype" + kwargs = {dtype_key: getattr(torch, args.dtype), "device_map": {"": "cpu"}, "output_loading_info": True} if args.flash_attention_2: kwargs["attn_implementation"] = "flash_attention_2" @@ -70,7 +85,9 @@ class FrozenQwenLocalEncoder: model.visual.to(self.device) model.audio_tower.to(self.device) self.model = model - self.processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id) + # Keep PIL/slow image preprocessing consistent with the original B1 + # definition, independently of newer Transformers' fast defaults. + self.processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id, use_fast=False) self.max_pixels = args.max_pixels self.image_batch_size = args.image_batch_size @@ -87,8 +104,9 @@ class FrozenQwenLocalEncoder: tokens = self.model.get_image_features(pixels, grid) merge = int(self.model.config.vision_config.spatial_merge_size) counts = (grid.prod(-1) // (merge * merge)).tolist() - if tokens.ndim != 2 or tokens.shape != (sum(counts), self.visual_dim): - raise RuntimeError(f"Unexpected image encoder shape: {tuple(tokens.shape)}") + if not isinstance(tokens, torch.Tensor) or tokens.shape != (sum(counts), self.visual_dim): + raise RuntimeError(f"Unexpected image encoder output: {type(tokens).__name__}, " + f"shape={getattr(tokens, 'shape', None)}") # Split BEFORE spatial mean: each person has one vote regardless of # box size / number of visual tokens. pieces.extend(part.float().mean(0).cpu() for part in tokens.split(counts)) @@ -108,8 +126,10 @@ class FrozenQwenLocalEncoder: mask = batch["attention_mask"].to(self.device) tokens = self.model.get_audio_features(features, feature_attention_mask=mask) _, lengths = self.model.audio_tower._get_feat_extract_output_lengths(mask.sum(-1)) - if tokens.ndim != 2 or tokens.shape != (int(lengths.sum()), self.audio_dim) or not len(tokens): - raise RuntimeError(f"Unexpected valid audio encoder shape: {tuple(tokens.shape)}") + if (not isinstance(tokens, torch.Tensor) + or tokens.shape != (int(lengths.sum()), self.audio_dim) or not len(tokens)): + raise RuntimeError(f"Unexpected valid audio encoder output: {type(tokens).__name__}, " + f"shape={getattr(tokens, 'shape', None)}") # get_audio_features removes padding and verifies the output length. return tokens.float().mean(0).cpu() @@ -371,6 +391,11 @@ def parse_args(argv=None): def main() -> None: args = parse_args() + if not args.preview_only: + import transformers + + # Fail before indexing/decoding source media, not after expensive I/O. + validate_local_transformers_version(transformers.__version__) names = [normalize_sample_name(value, args.dataset) for value in args.samples] if len(set(names)) != len(names): raise ValueError("Duplicate samples") @@ -417,7 +442,7 @@ def main() -> None: "config": { "encoder_id": args.encoder_id or args.model_id, "model_config_sha256": encoder.model_config_sha256, - "transformers": "4.52.3", "dtype": args.dtype, + "transformers": encoder.transformers_version, "dtype": args.dtype, "attention": "flash_attention_2" if args.flash_attention_2 else "default", "visual_dim": encoder.visual_dim, "audio_dim": encoder.audio_dim, "visual_fps": args.video_fps, "visual_window_seconds": args.visual_window_seconds, diff --git a/requirements.txt b/requirements.txt index 510e155..18313e9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,7 @@ torch numpy +# Default for Thinker replay. The B1 local extractor also supports 4.57.6; +# keep that version in existing GPU containers (see docs/local_feature_baseline.md). transformers==4.52.3 accelerate qwen-omni-utils[decord] diff --git a/test_local_turn_features.py b/test_local_turn_features.py index aab353e..c45b604 100644 --- a/test_local_turn_features.py +++ b/test_local_turn_features.py @@ -12,6 +12,7 @@ import sys import tempfile import unittest from types import SimpleNamespace +from unittest.mock import patch import numpy as np import torch @@ -23,6 +24,8 @@ from extract_turn_local_features import ( validate_source_timeline, model_config_digest, box_coordinate_offset, frame_crops, parse_args, preview_boxes, + TESTED_TRANSFORMERS_VERSIONS, validate_local_transformers_version, + main as extract_main, ) from local_turn_features import ( LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes, load_local_features, @@ -52,6 +55,13 @@ class LocalTests(unittest.TestCase): def setUp(self): torch.set_num_threads(1) + def test_tested_transformers_versions(self): + for version in ("4.52.3", "4.57.6"): + self.assertEqual(validate_local_transformers_version(version), version) + for version in ("4.50.0.dev0", "4.57.6.dev0", "5.0.0"): + with self.subTest(version=version), self.assertRaisesRegex(RuntimeError, "found " + version): + validate_local_transformers_version(version) + def test_half_open_grid_and_pooling(self): self.assertEqual(list(causal_frame_ticks(.2, 2)), [0]) self.assertEqual(list(causal_frame_ticks(1., 2)), [0, 1]) @@ -263,6 +273,67 @@ class LocalTests(unittest.TestCase): get_image_features=lambda *_: torch.tensor([[1., 1.], [3., 3.], [5., 5.]])) actual = encoder.images([object(), object()]) self.assertTrue(torch.equal(actual, torch.tensor([[1., 1.], [4., 4.]]))) + # A future API returning per-image tensors must not silently change + # the pooling rule or crash with an uninformative AttributeError. + encoder.model.get_image_features = lambda *_: (torch.ones(1, 2), torch.ones(2, 2)) + with self.assertRaisesRegex(RuntimeError, "Unexpected image encoder output"): + encoder.images([object(), object()]) + + def test_extraction_saves_runtime_version_and_rejects_mixed_versions(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for folder in ("features", "targets", "local"): + (root / folder).mkdir() + reference, side, labels = fixture() + torch.save(reference, root / "features/Session_4.pt") + np.save(root / "targets/Session_4.npy", labels) + media = SimpleNamespace(clips=[], metadata=lambda: {}) + argv = ["extract_turn_local_features.py", "--data-root", str(root), + "--feature-root", str(root / "features"), "--target-root", str(root / "targets"), + "--output-root", str(root / "local"), "--samples", "4", "--overwrite"] + contracts = [] + for version in TESTED_TRANSFORMERS_VERSIONS: + fake_encoder = SimpleNamespace(transformers_version=version, visual_dim=6, audio_dim=5, + model_config_sha256="same-checkpoint") + with patch.object(sys, "argv", argv), \ + patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__=version)}), \ + patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])), \ + patch("extract_turn_local_features.FrozenQwenLocalEncoder", return_value=fake_encoder), \ + patch("extract_turn_local_features.extract_local", return_value=copy.deepcopy(side)): + extract_main() + output = root / "local/Session_4.pt" + _, contract = load_local_features(output, reference, sample_id="Session_4", + dataset="easycom", modalities="av") + self.assertEqual(contract["config"]["transformers"], version) + contracts.append(contract) + with self.assertRaisesRegex(ValueError, "contract differs"): + load_local_features(output, reference, sample_id="Session_4", dataset="easycom", + modalities="av", expected_contract=contracts[0]) + + def test_unsupported_version_fails_before_media_and_preview_still_works(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + reference, _, labels = fixture() + torch.save(reference, root / "Session_4.pt") + np.save(root / "Session_4.npy", labels) + argv = ["extract_turn_local_features.py", "--data-root", str(root), + "--feature-root", str(root), "--target-root", str(root), + "--output-root", str(root / "preview"), "--samples", "4"] + with patch.object(sys, "argv", argv), \ + patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="5.0.0")}), \ + patch("extract_turn_local_features.build_media_and_boxes") as build: + with self.assertRaisesRegex(RuntimeError, "found 5.0.0"): + extract_main() + build.assert_not_called() + media = SimpleNamespace(clips=[], metadata=lambda: {}) + with patch.object(sys, "argv", argv + ["--preview-only"]), \ + patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="5.0.0")}), \ + patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])), \ + patch("extract_turn_local_features.preview_boxes") as preview, \ + patch("extract_turn_local_features.FrozenQwenLocalEncoder") as encoder: + extract_main() + preview.assert_called_once() + encoder.assert_not_called() def test_extraction_never_encodes_future_audio_and_saves_ragged_people(self): class Media: @@ -294,17 +365,19 @@ class LocalTests(unittest.TestCase): self.assertTrue(result["audio_valid"].all()) -@unittest.skipUnless(importlib.util.find_spec("transformers"), "optional: install transformers==4.52.3") +@unittest.skipUnless(importlib.util.find_spec("transformers"), "optional: install a tested Transformers version") class TinyQwenEncoderTests(unittest.TestCase): - def test_real_encoder_apis_without_downloading_weights(self): + def setUp(self): import transformers - if transformers.__version__ != "4.52.3": - self.skipTest("Pinned encoder API test requires transformers==4.52.3") + if transformers.__version__ not in TESTED_TRANSFORMERS_VERSIONS: + self.skipTest(f"Encoder API tests require one of {TESTED_TRANSFORMERS_VERSIONS}") + torch.set_num_threads(1) + + def tiny_model_and_processor(self): from transformers import ( Qwen2_5OmniThinkerConfig, Qwen2_5OmniThinkerForConditionalGeneration, Qwen2VLImageProcessor, WhisperFeatureExtractor, ) - torch.set_num_threads(1) config = Qwen2_5OmniThinkerConfig( audio_config=dict(d_model=16, encoder_layers=1, encoder_attention_heads=2, encoder_ffn_dim=32, output_dim=8), @@ -313,14 +386,18 @@ class TinyQwenEncoderTests(unittest.TestCase): text_config=dict(vocab_size=32, hidden_size=8, intermediate_size=16, num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2), ) + model = Qwen2_5OmniThinkerForConditionalGeneration(config).eval().requires_grad_(False) + processor = SimpleNamespace(image_processor=Qwen2VLImageProcessor(), + feature_extractor=WhisperFeatureExtractor(feature_size=128)) + return model, processor + + def test_real_encoder_apis_without_downloading_weights(self): encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder) - encoder.model = Qwen2_5OmniThinkerForConditionalGeneration(config).eval().requires_grad_(False) + encoder.model, encoder.processor = self.tiny_model_and_processor() encoder.model.model = None encoder.model.lm_head = None encoder.visual_dim = encoder.audio_dim = 8 encoder.device, encoder.max_pixels, encoder.image_batch_size = torch.device("cpu"), 224 * 224, 8 - encoder.processor = SimpleNamespace(image_processor=Qwen2VLImageProcessor(), - feature_extractor=WhisperFeatureExtractor(feature_size=128)) crops = [Image.new("RGB", (56, 84), "red"), Image.new("RGB", (112, 56), "blue")] batch = encoder.images(crops) single = encoder.images(crops[:1]) @@ -332,6 +409,44 @@ class TinyQwenEncoderTests(unittest.TestCase): self.assertTrue(torch.isfinite(audio).all()) self.assertFalse(audio.requires_grad) + @unittest.skipUnless(importlib.util.find_spec("accelerate"), "optional: install accelerate for checkpoint loading") + def test_constructor_loads_local_checkpoint_and_preserves_encoder_weights(self): + import transformers + + model, processor = self.tiny_model_and_processor() + expected = {key: value.clone() for key, value in model.state_dict().items() + if key.startswith(("visual.", "audio_tower."))} + with tempfile.TemporaryDirectory() as directory: + model.save_pretrained(directory) + args = SimpleNamespace(model_id=directory, device="cpu", dtype="float32", + flash_attention_2=False, max_pixels=224 * 224, image_batch_size=8) + with patch.object(transformers.Qwen2_5OmniProcessor, "from_pretrained", return_value=processor) as load: + encoder = FrozenQwenLocalEncoder(args) + load.assert_called_once_with(directory, use_fast=False) + self.assertEqual(encoder.transformers_version, transformers.__version__) + self.assertIsNone(encoder.model.model) + self.assertIsNone(encoder.model.lm_head) + self.assertFalse(encoder.model.training) + self.assertFalse(any(p.requires_grad for p in encoder.model.parameters())) + for key, value in encoder.model.state_dict().items(): + self.assertTrue(torch.equal(value, expected[key]), key) + self.assertEqual(encoder.images([Image.new("RGB", (56, 84), "red")]).shape, (1, 8)) + self.assertEqual(encoder.audio(np.zeros(3200, dtype=np.float32)).shape, (8,)) + + def test_constructor_rejects_incomplete_encoder_weights(self): + import transformers + + args = SimpleNamespace(model_id="unused", device="cpu", dtype="float32", flash_attention_2=False) + for info in ({"missing_keys": ["visual.patch_embed.proj.weight"]}, + {"missing_keys": ["audio_tower.conv1.weight"]}, + {"mismatched_keys": [("visual.weight", (2,), (3,))]}, + {"error_msgs": ["broken checkpoint"]}): + with self.subTest(info=info), \ + patch.object(transformers.Qwen2_5OmniThinkerForConditionalGeneration, + "from_pretrained", return_value=(None, info)), \ + self.assertRaisesRegex(RuntimeError, "incompletely loaded weights"): + FrozenQwenLocalEncoder(args) + class RoundTripTests(unittest.TestCase): def test_train_infer_easycom_and_ego4d(self):