Access Dataset Sample
Problem Statement
Given a PyTorch Dataset, extract a specific sample and return information about it.
Background
PyTorch Datasets support indexing — you can access individual samples, check the total number of samples, and iterate through the dataset.
Your Task
Write a function get_sample_info(dataset, index) that returns a dictionary with the total number of samples, the shape of the feature at the given index, and the label value at that index.
Output Format
Return a dictionary with keys: "num_samples" (int), "feature_shape" (list), "label_value" (int).
Example:
dataset with features [[1,2,3],[4,5,6]], labels [0,1], index=0
{"num_samples": 2, "feature_shape": [3], "label_value": 0}We use len(dataset) for num_samples, dataset[0][0].shape for feature shape, and dataset[0][1].item() for label
Constraints:
- Index will be valid (within dataset bounds)
- Return feature_shape as a list, not torch.Size
- label_value should be a Python int, not tensor
1. Background Knowledge
PyTorch Dataset and TensorDataset fundamentals:
- PyTorch Dataset is an abstract class implementing len() and getitem() for iterable data access
- TensorDataset combines multiple tensors into a dataset where dataset[i] returns a tuple (features[i], labels[i])
- Indexing returns tuples of tensors - critical for unpacking correctly
- len(dataset) returns the length of the shortest input tensor
Key tensor operations needed:
- tensor.shape → tuple of dimensions
- tensor.tolist() → Python native types
- int(tensor.item()) → scalar tensor to Python int
2. Algorithm Approach
Direct indexing with tuple unpacking:
sample = dataset[index] # Returns (features_tensor, label_tensor)
No complex algorithms needed - pure O(1) direct access leveraging PyTorch's optimized tensor storage.
Three independent computations:
- Dataset size: len(dataset) → O(1)
- Feature shape extraction: sample.shape → O(1)
- Label scalarization: int(sample.item()) → O(1)
3. Step-by-Step Strategy
def get_sample_info(dataset, index):
# Step 1: Get total samples
num_samples = len(dataset)
# Step 2: Extract sample tuple
sample = dataset[index] # (features, label)
# Step 3: Extract feature shape as list
feature_shape = list(sample.shape)
# Step 4: Convert label tensor to Python int
label_value = int(sample.item())
# Step 5: Return exact dictionary format
return {
"num_samples": num_samples,
"feature_shape": feature_shape,
"label_value": label_value
}
Execution flow:
- len(dataset) → integer count
- dataset[index] → tuple (features_tensor, label_tensor)
- features_tensor.shape → torch.Size → list() conversion
- label_tensor.item() → Python scalar → int() conversion
4. Common Pitfalls
| Mistake | Why It Fails | Fix |
|---|---|---|
| return tensor.shape | Returns torch.Size, not list | list(tensor.shape) |
| return sample | Returns tensor, not Python int | int(sample.item()) |
| features, labels = dataset[index] | Unpacks tuple correctly but shadows variable names | Use sample = dataset[index] then sample, sample |
| Missing .item() on scalar tensor | TypeError or tensor returned | Always use .item() for scalar extraction |
| sample.size() | Returns torch.Size | Use .shape property |
Type conversion cheatsheet:
tensor.shape → torch.Size() → list(tensor.shape) →
scalar_tensor → tensor() → tensor.item() → 0.0
scalar_tensor → tensor() → int(tensor.item()) → 0
5. Time & Space Complexity
Time Complexity: O(1)
- len(dataset): O(1) - stored metadata access
- dataset[index]: O(1) - direct tensor slicing
- Shape extraction: O(1) - metadata read
- .item(): O(1) - scalar extraction
Space Complexity: O(1)
- Only extracts one sample regardless of dataset size N
- Temporary tuple and metadata only
- Output dictionary: fixed 3 entries
Scalability:
Time: Independent of N (dataset size)
Space: Independent of N
Complete Working Solution
import torch
from torch.utils.data import TensorDataset
def get_sample_info(dataset, index):
"""Extract dataset info at specific index."""
num_samples = len(dataset)
sample = dataset[index]
return {
"num_samples": num_samples,
"feature_shape": list(sample.shape),
"label_value": int(sample.item())
}
# Test
features = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
labels = torch.tensor([0, 1])
dataset = TensorDataset(features, labels)
print(get_sample_info(dataset, 0))
# {"num_samples": 2, "feature_shape":, "label_value": 0}
Key insight: This problem tests understanding of PyTorch Dataset indexing protocol and tensor-to-Python type conversion - fundamental for any PyTorch data pipeline.