{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Export data for ML training\n", "\n", "Convert Pixeltable data to PyTorch DataLoader format for model training." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Problem\n", "\n", "You have prepared training data—images with labels, text with embeddings, or multimodal data—and need to export it for PyTorch model training.\n", "\n", "| Data type | Use case |\n", "|-----------|----------|\n", "| Images + labels | Image classification |\n", "| Text + embeddings | Fine-tuning embeddings |\n", "| Audio + transcripts | Speech model training |" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Solution\n", "\n", "**What's in this recipe:**\n", "\n", "- Convert query results to PyTorch Dataset\n", "- Use with DataLoader for batch training\n", "- Export to Parquet for external tools\n", "\n", "You use `query.to_pytorch_dataset()` to create an iterable dataset compatible with PyTorch DataLoader." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Setup" ] }, { "cell_type": "code", "metadata": {}, "source": [ "%pip install -qU pixeltable torch torchvision" ], "execution_count": null, "outputs": [] }, { "cell_type": "code", "metadata": {}, "source": [ "import pixeltable as pxt\n", "import torch\n", "from torch.utils.data import DataLoader\n", "\n", "# Create a fresh directory\n", "pxt.drop_dir('pytorch_demo', force=True)\n", "pxt.create_dir('pytorch_demo')" ], "execution_count": 3, "outputs": [ { "output_type": "stream", "text": [ "Connected to Pixeltable database at: postgresql+psycopg://postgres:@/pixeltable?host=/Users/pjlb/.pixeltable/pgdata\n", "Created directory 'pytorch_demo'.\n" ] }, { "output_type": "execute_result", "data": { "text/plain": [ "" ] } } ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Create sample training data" ] }, { "cell_type": "code", "metadata": {}, "source": [ "# Create table with images and labels\n", "training_data = pxt.create_table(\n", " 'pytorch_demo/training_data', {'image': pxt.Image, 'label': pxt.Int}\n", ")" ], "execution_count": 4, "outputs": [ { "output_type": "stream", "text": [ "Created table 'training_data'.\n" ] } ] }, { "cell_type": "code", "metadata": {}, "source": [ "# Insert sample images with labels\n", "base_url = 'https://raw.githubusercontent.com/pixeltable/pixeltable/main/docs/resources/images'\n", "samples = [\n", " {'image': f'{base_url}/000000000036.jpg', 'label': 0}, # cat\n", " {'image': f'{base_url}/000000000090.jpg', 'label': 1}, # other\n", " {'image': f'{base_url}/000000000139.jpg', 'label': 1}, # other\n", "]\n", "training_data.insert(samples)" ], "execution_count": 5, "outputs": [ { "output_type": "stream", "text": [ "Inserting rows into `training_data`: 3 rows [00:00, 659.03 rows/s]\n", "Inserted 3 rows with 0 errors.\n" ] }, { "output_type": "execute_result", "data": { "text/plain": [ "3 rows inserted, 6 values computed." ] } } ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Export to PyTorch dataset" ] }, { "cell_type": "code", "metadata": {}, "source": [ "# Add a resize step to ensure all images have the same size\n", "training_data.add_computed_column(\n", " image_resized=training_data.image.resize((224, 224))\n", ")\n", "\n", "# Convert to PyTorch dataset\n", "# 'pt' format returns images as CxHxW tensors with values in [0,1]\n", "pytorch_dataset = training_data.select(\n", " training_data.image_resized, training_data.label\n", ").to_pytorch_dataset(image_format='pt')" ], "execution_count": 6, "outputs": [ { "output_type": "stream", "text": [ "Added 3 column values with 0 errors.\n" ] } ] }, { "cell_type": "code", "metadata": {}, "source": [ "# Use with PyTorch DataLoader\n", "dataloader = DataLoader(pytorch_dataset, batch_size=2)\n", "\n", "# Get first batch to verify the shape\n", "batch = next(iter(dataloader))\n", "batch[\n", " 'image_resized'\n", "].shape # Should be (2, 3, 224, 224) - batch_size x channels x height x width" ], "execution_count": 7, "outputs": [ { "output_type": "execute_result", "data": { "text/plain": [ "torch.Size([2, 3, 224, 224])" ] } } ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Export to Parquet for external tools" ] }, { "cell_type": "code", "metadata": {}, "source": [ "import tempfile\n", "from pathlib import Path\n", "\n", "# Export to Parquet for use with other ML tools\n", "export_path = Path(tempfile.mkdtemp()) / 'training_data'\n", "\n", "pxt.io.export_parquet(\n", " training_data.select(training_data.label), # Non-image columns\n", " parquet_path=export_path,\n", ")" ], "execution_count": 8, "outputs": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Explanation\n", "\n", "**Export methods:**\n", "\n", "| Method | Output | Use case |\n", "|--------|--------|----------|\n", "| `to_pytorch_dataset()` | PyTorch IterableDataset | Direct training |\n", "| `export_csv()` | CSV files | Tabular data, spreadsheets |\n", "| `export_json()` | JSON files | Structured data, APIs |\n", "| `export_parquet()` | Parquet files | External tools, sharing |\n", "| `export_lancedb()` | LanceDB | Vector search apps |\n", "| `to_coco_dataset()` | COCO JSON | Object detection |\n", "\n", "**Image format options:**\n", "\n", "| Format | Shape | Values | Use |\n", "|--------|-------|--------|-----|\n", "| `'pt'` | CxHxW | [0, 1] float32 | PyTorch models |\n", "| `'np'` | HxWxC | [0, 255] uint8 | NumPy processing |\n", "\n", "**DataLoader tips:**\n", "\n", "- Data is cached to disk for efficient repeated loading\n", "- Use `num_workers > 0` for parallel data loading\n", "- Filter/transform data before export to reduce size" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## See also\n", "\n", "- [Sample data for training](https://docs.pixeltable.com/howto/cookbooks/data/data-sampling) - Stratified sampling\n", "- [Import Parquet files](https://docs.pixeltable.com/howto/cookbooks/data/data-import-parquet) - Parquet import/export" ] } ], "metadata": { "kernelspec": { "display_name": "pixeltable", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.11" } }, "nbformat": 4, "nbformat_minor": 2 }