torchvision-vibecoding-project/todo.md

Project To-Do List

This list outlines the steps required to complete the Torchvision Finetuning project, derived from prompt_plan.md.

Phase 1: Foundation & Setup

• Set up project structure (directories: configs, data, models, utils, tests, scripts).
• Initialize Git repository.
• Create .gitignore file (ignore data, outputs, logs, .venv, caches, *.pth).
• Initialize pyproject.toml using uv init, set Python 3.10.
• Add core dependencies (torch, torchvision, ruff, numpy, Pillow, pytest) using uv add.
• Create pre-commit-config.yaml and configure ruff hooks (format, lint, import sort).
• Create __init__.py files in necessary directories.
• Create empty placeholder files (train.py, test.py, configs/base_config.py, utils/data_utils.py, models/detection.py, tests/conftest.py).
• Create basic README.md.
• Install pre-commit hooks (pre-commit install).
• Create scripts/download_data.sh script.
  • Check if data exists.
  • Create data/ directory.
  • Use wget to download PennFudanPed dataset.
  • Use unzip to extract data.
  • Remove zip file after extraction.
  • Add informative print messages.
  • Make script executable (chmod +x).
• Ensure .gitignore ignores data/.
• Implement base configuration in configs/base_config.py (base_config dictionary).
• Implement specific experiment configuration in configs/pennfudan_maskrcnn_config.py (config dictionary, importing/updating base config).

Phase 2: Data Handling & Model

• Implement PennFudanDataset class in utils/data_utils.py.
  • __init__: Load image and mask paths.
  • __getitem__: Load image/mask, parse masks, generate targets (boxes, labels, masks, image_id, area, iscrowd), apply transforms.
  • __len__: Return dataset size.
• Implement get_transform(train) function in utils/data_utils.py (using torchvision.transforms.v2).
• Implement collate_fn(batch) function in utils/data_utils.py.
• Implement get_maskrcnn_model(num_classes, ...) function in models/detection.py.
  • Load pre-trained Mask R-CNN (maskrcnn_resnet50_fpn_v2).
  • Replace box predictor head (FastRCNNPredictor).
  • Replace mask predictor head (MaskRCNNPredictor).

Phase 3: Training Script & Core Logic

• Set up basic train.py structure.
  • Add imports.
  • Implement argparse for --config argument.
  • Implement dynamic config loading (importlib).
  • Set random seeds.
  • Determine compute device (cuda or cpu).
  • Create output directory structure (outputs/<config_name>/checkpoints).
  • Instantiate PennFudanDataset (train).
  • Instantiate DataLoader (train) using collate_fn.
  • Instantiate model using get_maskrcnn_model.
  • Move model to device.
  • Add if __name__ == "__main__": guard.
• Implement minimal training step in train.py.
  • Instantiate optimizer (torch.optim.SGD).
  • Set model.train().
  • Fetch one batch.
  • Move data to device.
  • Perform forward pass (loss_dict = model(...)).
  • Calculate total loss (sum(...)).
  • Perform backward pass (optimizer.zero_grad(), loss.backward(), optimizer.step()).
  • Print/log loss for the single step (and temporarily exit).
• Implement logging setup in utils/log_utils.py (setup_logging function).
  • Configure logging.basicConfig for file and console output.
• Integrate logging into train.py.
  • Call setup_logging.
  • Replace print with logging.info.
  • Log config, device, and training progress/losses.
• Implement full training loop in train.py.
  • Remove single-step exit.
  • Add LR scheduler (torch.optim.lr_scheduler.StepLR).
  • Add epoch loop.
  • Add batch loop, integrating the single training step logic.
  • Log loss periodically within the batch loop.
  • Step the LR scheduler at the end of each epoch.
  • Log total training time.
• Implement checkpointing in train.py.
  • Define checkpoint directory.
  • Implement logic to find and load the latest checkpoint (resume training).
  • Save checkpoints periodically (based on frequency or final epoch).
    • Include epoch, model state, optimizer state, scheduler state, config.
  • Log checkpoint loading/saving.

Phase 4: Evaluation & Testing

• Add evaluation dependencies (pycocotools - optional initially).
• Create utils/eval_utils.py and implement evaluate function.
  • Set model.eval().
  • Use torch.no_grad().
  • Loop through validation/test dataloader.
  • Perform forward pass.
  • Calculate/aggregate metrics (start with average loss, potentially add mAP later).
  • Log evaluation metrics and time.
  • Return metrics.
• Integrate evaluation into train.py.
  • Create validation Dataset and DataLoader (using torch.utils.data.Subset).
  • Call evaluate at the end of each epoch.
  • Log validation metrics.
  • (Later) Implement logic to save the best model based on validation metric.
• Implement test.py script.
  • Reuse argument parsing, config loading, device setup, dataset/dataloader (test split), model creation from train.py.
  • Add --checkpoint argument.
  • Load model weights from the specified checkpoint.
  • Call evaluate function using the test dataloader.
  • Log/print final evaluation results.
  • Setup logging for testing (e.g., test.log).
• Create unit tests in tests/ using pytest.
  • tests/test_config.py: Test config loading.
  • tests/test_model.py: Test model creation and head configuration.
  • tests/test_data_utils.py: Test dataset instantiation, length, and item format (requires data).
  • (Optional) Use fixtures in tests/conftest.py if needed.
• Add pytest execution to pre-commit-config.yaml.
• Test pre-commit hooks (pre-commit run --all-files).

Phase 5: Refinement & Documentation

• Refine error handling in train.py and test.py (try...except).
• Add configuration validation checks.
• Improve evaluation metrics (e.g., implement mAP in evaluate function).
• Add more data augmentations to get_transform(train=True).
• Expand README.md significantly.
  • Goals
  • Detailed Setup
  • Configuration explanation
  • Training instructions (including resuming)
  • Testing instructions
  • Project Structure overview
  • Dependencies list
  • (Optional) Results section
• Perform final code quality checks (ruff format ., ruff check . --fix).
• Ensure all pre-commit hooks pass.