Add Tensor Input Support: Enable .pt file processing with <tensor> tags for latent representations #6504
Conversation
This commit fixes two critical bugs that would prevent tensor input functionality from working: 1. Missing tensor_idx initialization in StdTemplateInputs.__post_init__ - Added tensor_idx = 0 initialization (line 122) - Without this, the code would fail when trying to access tensor_idx during template processing 2. Missing tensor list handling in StdTemplateInputs.__post_init__ - Added check to convert tensors to list if not already one (lines 131-132) - This ensures consistency with how images, videos, and audios are handled 3. Missing tensor support in dataset preprocessor - Added 'tensors' to RowPreprocessor.standard_keys (line 24) - Added tensors_keys mapping for 'tensors' and 'tensor' (line 39) - Added 'tensors' to mm_type loop (line 40) - This is critical for dataset loading to work with tensor files These fixes ensure that tensor inputs are properly initialized and handled throughout the template processing pipeline, matching the existing patterns for images, videos, and audios.
This commit adds: 1. Unit Tests (tests/test_align/test_template/test_tensor.py): - test_load_tensor: Tests basic tensor loading from .pt files - test_load_batch_tensor: Tests batch tensor loading (4D tensors) - test_tensor_to_images_single: Tests single tensor to PIL Image conversion - test_tensor_to_images_batch: Tests batch tensor to PIL Images conversion - test_tensor_to_images_grayscale: Tests grayscale tensor conversion - test_tensor_to_images_2d: Tests 2D tensor conversion - test_tensor_normalization: Tests proper value normalization [0,1] -> [0,255] - test_template_inputs_with_tensors: Tests StdTemplateInputs tensor support - test_infer_request_with_tensors: Tests InferRequest tensor support - test_mixed_media: Tests tensors with images, videos, audios - test_dataset_preprocessor_tensor_support: Tests dataset preprocessor - test_special_tokens: Tests tensor special tokens registration Total: 12 comprehensive test cases covering all tensor functionality 2. Documentation (docs/source_en/Customization/Tensor-support.md): - Complete guide to tensor input support - Supported tensor formats and shapes - Usage examples (basic, multiple tensors, mixed media) - Dataset format examples for all training modes - Command line usage - Creating tensor files (with medical imaging example) - Custom template implementation - Technical details and troubleshooting - Use cases: medical report generation, scientific data analysis 3. Documentation Updates (docs/source_en/Customization/Custom-dataset.md): - Added tensors to standard dataset keys documentation - Added <tensor> tag to multimodal section - Added tensor examples for pre-training and supervised fine-tuning - Added reference link to detailed tensor documentation The tests follow the existing test patterns in ms-swift and cover: - Tensor loading (single, batch) - Tensor-to-image conversion (RGB, grayscale, 2D, 3D, 4D) - Template inputs integration - Dataset preprocessor integration - Mixed multimodal usage The documentation provides comprehensive guidance for users wanting to: - Train models with tensor inputs - Use tensors for medical imaging, scientific data, or custom features - Integrate tensors with existing multimodal pipelines
Summary of ChangesHello @Marshall-mk, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for tensor inputs (.pt files), which is a great new feature for handling specialized data like medical scans. The implementation is thorough, covering data preprocessing, template handling, and tensor-to-image conversion logic. The addition of comprehensive documentation and a full suite of unit tests is excellent and greatly improves the quality of the contribution.
My review includes one suggestion for improving code maintainability by removing redundant imports. Overall, this is a high-quality pull request.
| from PIL import Image | ||
| import numpy as np |
There was a problem hiding this comment.
The import from PIL import Image is redundant, as Image is already imported at the module level (line 20). Additionally, import numpy as np is duplicated in the _single_tensor_to_image method. To improve maintainability, you could remove the redundant Image import. For numpy, you could import it once at the top of the file and remove the local imports.
| from PIL import Image | ||
| import numpy as np |
There was a problem hiding this comment.
The import from PIL import Image is redundant as Image is already imported at the module level (line 20). Additionally, import numpy as np is duplicated in the _tensor_to_images method. To improve maintainability, you could remove the redundant Image import. For numpy, you could import it once at the top of the file and remove the local imports.
|
Hello. Thanks for your PR. Could you tell me what additional advantages using |
The key difference here is that you can't load tensors (.pt files) using the You might be wondering why we’d want to load tensors at all. In my case, I work with medical images such as CT or MRI scans, which often contain multiple slices and have large dimensions. These images make training VLMs very compute-intensive. To efficiently process these images, we need to compress them using an encoder, save the resulting embeddings as tensors, and then use the VLMs available in ms-swift. Now, a naive approach would be to convert the embeddings (the tensors) into videos (from .pt files to .mp4), but that introduces the need for additional normalizations and scaling, which actually leads to some disparity in the actual embeddings. This is why using directly is a better approach, as it avoids the unnecessary conversion and preserves the embeddings. |
3 participants
PR type
PR information
Write the detail information belongs to this PR.
Experiment results
Paste your experiment result here(if needed).