Add FiftyOneTorchDataset #5321
Add FiftyOneTorchDataset #5321
Conversation
…ther than view. Way faster.
WalkthroughThis pull request introduces a comprehensive set of tools and scripts for training machine learning models using the MNIST dataset with FiftyOne and PyTorch. The changes include a training script ( Changes
Possibly related PRs
Suggested reviewers
Poem
Tip CodeRabbit's docstrings feature is now available as part of our Early Access Program! Simply use the command Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (8)
fiftyone/utils/torch.py (1)
1605-1605: Useisfor None checksThis line compares to
Noneusing==. For clarity and correctness, preferis None.- if self._samples == None: + if self._samples is None:🧰 Tools
🪛 Ruff (0.8.2)
1605-1605: Comparison to
Noneshould becond is NoneReplace with
cond is None(E711)
docs/source/recipes/torch-dataset-examples/mnist_training.py (1)
12-39: [Function main: handle partial device availability?]Great job customizing device usage. Consider adding a fallback for cases where
"cuda:0"or the specified device is not available, e.g. defaulting to CPU.docs/source/recipes/torch-dataset-examples/utils.py (3)
1-1: Remove unused import if truly unnecessaryIt appears
import fiftyone as fois unused here. Removing it can improve clarity and avoid confusion.- import fiftyone as fo + # remove if truly not needed🧰 Tools
🪛 Ruff (0.8.2)
1-1:
fiftyoneimported but unusedRemove unused import:
fiftyone(F401)
131-131: Replace unnecessary ternary with direct comparisonUse a direct boolean expression to simplify logic.
- shuffle = True if split_tag == "train" else False + shuffle = (split_tag == "train")🧰 Tools
🪛 Ruff (0.8.2)
131-131: Remove unnecessary
True if ... else FalseRemove unnecessary
True if ... else False(SIM210)
142-150: [setup_model: consider user flexibility]Some users might need a different backbone or weights. Possibly add a parameter to switch from ResNet18 to other architectures?
docs/source/recipes/torch-dataset-examples/simple_training_example.ipynb (1)
68-68: Typographical correction in markdown"traing" => "training" to maintain clarity in documentation.
- Now we will look at an actual traing script with `FiftyOneTorchDataset` + Now we will look at an actual training script with `FiftyOneTorchDataset`fiftyone/core/collections.py (2)
10973-10978: Consider adding type hints for better IDE supportThe method could benefit from type hints for the
get_itemparameter and return type to improve IDE support and code clarity.- def to_torch(self, get_item, **kwargs): + def to_torch(self, get_item: Callable[[Any], Any], **kwargs) -> "FiftyOneTorchDataset":
10973-10978: Consider adding validation for theget_itemparameterThe method should validate that
get_itemis a callable to provide a better error message.def to_torch(self, get_item, **kwargs): + if not callable(get_item): + raise ValueError( + "The 'get_item' parameter must be a callable that accepts a sample and " + "returns the corresponding tensor(s)" + ) from fiftyone.utils.torch import FiftyOneTorchDataset return FiftyOneTorchDataset(self, get_item, **kwargs)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
docs/source/recipes/torch-dataset-examples/mnist_training.py(1 hunks)docs/source/recipes/torch-dataset-examples/simple_training_example.ipynb(1 hunks)docs/source/recipes/torch-dataset-examples/utils.py(1 hunks)fiftyone/core/collections.py(1 hunks)fiftyone/utils/torch.py(4 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
docs/source/recipes/torch-dataset-examples/utils.py
1-1: fiftyone imported but unused
Remove unused import: fiftyone
(F401)
131-131: Remove unnecessary True if ... else False
Remove unnecessary True if ... else False
(SIM210)
fiftyone/utils/torch.py
1605-1605: Comparison to None should be cond is None
Replace with cond is None
(E711)
🔇 Additional comments (25)
fiftyone/utils/torch.py (10)
13-15: [Imports look good]
No issues found with these imports.
32-34: [Imports look good]
No issues found with these imports.
41-41: [Distributed import looks good]
No issues found with this import.
1470-1551: [Comprehensive class docstring, few content suggestions]
The docstring is highly informative. However, consider briefly specifying the expected return type in get_item (e.g., dict, tensor, or custom object) to clarify type-checking needs.
1610-1614: [Worker init logic seems correct]
The worker_init approach is aligned with best practices for DataLoader workers. Ensure thorough testing in multi-worker scenarios to confirm memory usage remains stable.
1640-1645: [Good specialized logic for cached fields]
This dictionary-based approach for returning sample data is solid for decoupling data loading from the main process. Continue ensuring the data schema remains stable across training epochs.
2220-2244: [NumpySerializedList Implementation]
Serializing data to a numpy array of bytes is a clever approach. Be aware of potential overhead for extremely large datasets.
2255-2266: [TorchSerializedList Implementation]
Converting to torch tensors is convenient for consistency with PyTorch. Looks solid.
2268-2298: [Shared memory approach]
The design for distributing serialized data across local DDP ranks is well-thought-out. Ensure proper cleanup of shared memory resources after training completes to prevent memory leaks.
2300-2351: [DDP Communication Routines]
These distributed utilities (e.g., all_gather, local_scatter) are essential for multi-rank synchronization. Make sure you test corner cases with a single rank or large multi-GPU clusters.
docs/source/recipes/torch-dataset-examples/mnist_training.py (3)
12-39: [Potential File Save Race Condition]
When saving models (torch.save(...)), measuring the overhead or ensuring concurrency safety might be needed if multiple processes attempt to save.
62-83: [Function train_epoch: structured approach looks good]
Make sure large-scale loads into batch don’t cause memory fragmentation under certain expansions of dataset size or dynamic input shapes.
86-127: [Function validation: good usage of no_grad()]
This function properly detaches computations. Just be mindful that updates to dataset._dataset.select(...) inside your loop might cause overhead if repeated too frequently on large sets.
docs/source/recipes/torch-dataset-examples/utils.py (10)
17-37: [Function get_item_quickstart: correct bounding box transformation]
Looks good. Keep in mind that large bounding box lists might impact performance if done repeatedly.
39-59: [Function get_item_cached_quickstart: consistent approach with get_item_quickstart]
The code is consistent. Validate that the fields in sample_dict exactly match your bounding box transformations to avoid index or key errors.
61-63: [Collate function: flexible approach]
Returning tuples of (image, label) as separate items is typical. If you foresee multi-output tasks or dynamic batching, consider a more advanced approach.
65-75: [Simple DataLoader creation: good reference]
This sets the stage for an easy integration with FiftyOneTorchDataset. No issues noted.
77-84: [ids_in_dataloader: demonstration function]
This demonstration function is helpful for basic usage. Keep an eye out for large memory usage if the dataset is big.
86-101: [Mapping MNIST indexes to strings: straightforward & clear]
No issues found. This is a good reference for labeling in logs.
103-107: [convert_and_normalize: future deprecation]
The warning about ToTensor() is relevant. Keep in mind that a future replacement with v2.ToImage() and v2.ToDtype(...) might be necessary.
110-118: [mnist_get_item: looks correct]
The function is straightforward. If the label format changes, ensure this integer parsing remains valid.
120-139: [create_dataloaders: flexible for train/val/test splits]
Ensure you handle the scenario when one or more splits are empty or missing tags.
🧰 Tools
🪛 Ruff (0.8.2)
131-131: Remove unnecessary True if ... else False
Remove unnecessary True if ... else False
(SIM210)
152-155: [setup_optim: standard setup]
This is a reasonable default. If hyperparameter tuning is frequent, consider making LR or L2 configurable at runtime.
fiftyone/core/collections.py (2)
10973-10978: LGTM: Clean implementation of the to_torch method
The implementation correctly delegates to FiftyOneTorchDataset while maintaining a clean interface. The docstring reference to the implementation class provides good discoverability.
10973-10978: Verify PyTorch dependency availability
Since this method integrates with PyTorch, we should verify that PyTorch is available when this method is called.
✅ Verification successful
PyTorch dependency is properly handled
The codebase already has proper PyTorch dependency verification through fou.ensure_torch() which is called before importing torch in the fiftyone/utils/torch.py module. This function verifies that both torch and torchvision are installed and importable. Since the to_torch() method imports from fiftyone.utils.torch, which handles the dependency check, no additional verification is needed in the method itself.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check if torch is listed as a dependency
grep -r "torch" setup.py requirements*.txt
Length of output: 41
Script:
#!/bin/bash
# Search for torch-related imports and try-except blocks
rg "torch" -A 5 -B 5
Length of output: 82613
Script:
#!/bin/bash
# Search for optional dependencies setup
rg "extras_require" setup.py
Length of output: 28
Script:
#!/bin/bash
# Look for any torch-related error handling
ast-grep --pattern 'try {
$$$
import torch
$$$
} catch ($$$) {
$$$
}'
Length of output: 79
Script:
#!/bin/bash
# Look for torch import handling in the specific file
rg "torch" -B 5 -A 5 fiftyone/core/collections.py
# Look for any import error handling patterns
rg "ImportError" -B 5 -A 5 fiftyone/core/collections.py
Length of output: 3228
Script:
#!/bin/bash
# Check the torch utility module for any import error handling
rg "import.*torch" -B 5 -A 5 fiftyone/utils/torch.py
# Look for any torch-related imports in __init__.py files
rg "import.*torch" -B 5 -A 5 "**/__init__.py"
Length of output: 2442
Script:
#!/bin/bash
# Look for ensure_torch function implementation
rg "ensure_torch" -B 5 -A 5 fiftyone/utils/
# Look for any torch-related imports in utils directory
rg "import.*torch" -B 5 -A 5 fiftyone/utils/
Length of output: 15094
Script:
#!/bin/bash
# Look for ensure_torch implementation in core/utils.py
rg "def ensure_torch" -B 5 -A 10 fiftyone/core/utils.py
Length of output: 488
|
MAGE training successful. Very coarse test, but nice to see still. |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (8)
docs/source/recipes/torch-dataset-examples/ddp_train.py (5)
39-45: Confirm dataset loading approachEach process is synchronously loading the same dataset locally, which can increase memory usage for large datasets. Confirm whether this is the intended design or if a shared or streamed approach would be more efficient.
81-81: Fix minor spelling in log message"New best lost achieved" appears to be a typo. Change "lost" to "loss" for clarity.
- print(f"New best lost achieved : {best_loss}. Saving model...") + print(f"New best loss achieved : {best_loss}. Saving model...")
86-90: Remove redundant.to(DEVICES[local_rank])callLine 89 and 90 both move the model to the same device. You can remove one to avoid duplication.
model = utils.setup_ddp_model( num_classes=num_classes, weights_path=f"{save_dir}/epoch_{best_epoch}.pt", ).to(DEVICES[local_rank]) - model.to(DEVICES[local_rank]) ddp_model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[DEVICES[local_rank]] )
125-125: Correct the spelling of “cummulative_loss”Use “cumulative_loss” for clearer communication.
- cummulative_loss = 0 + cumulative_loss = 0And update its usage accordingly in lines 142 and 157 as well.
Also applies to: 142-142, 157-157
143-144: Refactor nested if statementsYou can combine these nested if statements for simpler logic, per the static analysis suggestion.
-if local_rank == 0: - if batch_num % 100 == 0: +if local_rank == 0 and batch_num % 100 == 0: pbar.set_description(...)Also applies to: 193-194
🧰 Tools
🪛 Ruff (0.8.2)
143-144: Use a single
ifstatement instead of nestedifstatementsCombine
ifstatements usingand(SIM102)
docs/source/recipes/torch-dataset-examples/utils.py (3)
1-1: Remove unused import [fiftyone].The module is never referenced in the code, so you can safely delete this import line.
- import fiftyone as fo🧰 Tools
🪛 Ruff (0.8.2)
1-1:
fiftyoneimported but unusedRemove unused import:
fiftyone(F401)
118-137: Function create_dataloaders: simplify conditional assignment to boolean.Line 129 uses
shuffle = True if split_tag == "train" else False. This can be simplified to:shuffle = (split_tag == "train")🧰 Tools
🪛 Ruff (0.8.2)
129-129: Remove unnecessary
True if ... else FalseRemove unnecessary
True if ... else False(SIM210)
164-186: Function create_dataloaders_ddp: simplify conditional assignment to boolean.Line 175 uses
shuffle = True if split_tag == "train" else False. This can be simplified to:shuffle = (split_tag == "train")🧰 Tools
🪛 Ruff (0.8.2)
175-175: Remove unnecessary
True if ... else FalseRemove unnecessary
True if ... else False(SIM210)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
docs/source/recipes/torch-dataset-examples/ddp_train.py(1 hunks)docs/source/recipes/torch-dataset-examples/utils.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
docs/source/recipes/torch-dataset-examples/ddp_train.py
143-144: Use a single if statement instead of nested if statements
Combine if statements using and
(SIM102)
193-194: Use a single if statement instead of nested if statements
Combine if statements using and
(SIM102)
docs/source/recipes/torch-dataset-examples/utils.py
1-1: fiftyone imported but unused
Remove unused import: fiftyone
(F401)
129-129: Remove unnecessary True if ... else False
Remove unnecessary True if ... else False
(SIM210)
175-175: Remove unnecessary True if ... else False
Remove unnecessary True if ... else False
(SIM210)
🔇 Additional comments (17)
docs/source/recipes/torch-dataset-examples/ddp_train.py (3)
14-30: Ensure environment variables are defined
The code relies on environment variables like WORLD_SIZE, LOCAL_WORLD_SIZE, and RANK. Consider adding error handling or clear messages if they are not set, to prevent runtime failures in distributed settings.
176-190: Evaluate performance impact of per-batch dataset updates
Calling dataset._dataset.select(batch["id"]).set_values(...) and samples.save() for each batch can be costly for large datasets. Consider batching result writing or deferring it until after validation to reduce overhead.
[performance]
201-207: Clarify execution instructions
The inline comment shows an example command for 6 processes (--nproc-per-node=6), but the example sets --devices 2 3 4 5 6 7, which are 6 GPUs. It would be helpful to clarify these arguments to ensure they match real-world usage scenarios (e.g., confirming GPU device IDs and process counts).
docs/source/recipes/torch-dataset-examples/utils.py (14)
2-2: Import usage .
from fiftyone.utils.torch import FiftyOneTorchDataset is used below, so this import is necessary.
4-15: Imports and augmentation pipeline .
No immediate concerns or suggestions. The augmentation pipeline is clearly defined and easy to follow.
17-37: Function get_item_quickstart .
The workflow to open the image, compute bounding boxes, and apply augmentations is clear and logical. Your handling of empty detections with a zero-tensor is appropriately done.
39-59: Function get_item_cached_quickstart .
Similar to get_item_quickstart, this approach for working with a cached sample dictionary is clear and logically consistent.
61-62: Function simple_collate_fn .
This is a valid approach that returns data in an easily unpackable tuple. No changes needed.
65-75: Function create_dataloader_simple .
DataLoader creation is straightforward, and the use of FiftyOneTorchDataset.worker_init is correctly applied. No issues found.
77-84: Function ids_in_dataloader .
Retrieving IDs from batches for verification is a helpful debugging utility. The assertion for batch size ensures the logic is working as intended.
87-101: Function mnist_index_to_label_string .
Provides a neat mapping from label indices to strings. No concerns here.
103-105: Transformation convert_and_normalize .
It neatly converts images to the proper format for PyTorch and normalizes pixel values. All good here.
108-116: Function mnist_get_item .
Processing MNIST samples into a dictionary is well-executed. The label extraction is direct and understandable.
140-148: Function setup_model .
Configuring a ResNet18 plus linear head is appropriately done. Loading weights conditionally is also clear.
150-153: Function setup_optim .
Simple, well-defined SGD optimizer creation. No issues spotted.
158-162: Function setup_ddp_model .
Applying convert_sync_batchnorm is the recommended practice for DDP. Steps are logical.
188-190: No-op main guard .
The if __name__ == "__main__": pass block can be helpful when multiprocessing in notebooks. It’s fine to keep as-is.
What changes are proposed in this pull request?
FiftyOneTorchDataset.to_torchmethod.How is this patch tested? If it is not, please explain why.
Release Notes
Is this a user-facing change that should be mentioned in the release notes?
notes for FiftyOne users.
New
to_torchmethod allows for easy transfer to a fully functionaltorch.utils.data.Datasetobject.Please see recipes for examples and tutorials.
(Details in 1-2 sentences. You can just refer to another PR with a description
if this PR is part of a larger change.)
What areas of FiftyOne does this PR affect?
fiftyonePython library changesfiftyone.utils.torchSummary by CodeRabbit
New Features
Bug Fixes
Documentation