Lesson · 40 min · Free
Custom Dataset Training in PyTorch
Custom Dataset Training in PyTorch body { font-family: sans-serif; line-height: 1.6; margin: 20px; } h1, h2 { color: #2c3e50; } pre { background-color: #ecf0f1; padding: 15px; border-radius: 5px; overflow-x: auto; font-f
Custom Dataset Training in PyTorch
Welcome to this lesson on custom dataset training in PyTorch, a crucial skill for anyone working with novel biological data in machine learning applications. While pre-built datasets like ImageNet or MNIST are excellent for learning fundamentals, real-world problems in structural biology and drug discovery often involve unique data formats, varying sizes, and specialized annotations. PyTorch provides a flexible and powerful framework for handling such custom datasets, allowing you to seamlessly integrate your biological data into deep learning models. At an upper-undergraduate level, you should already be familiar with basic PyTorch tensors, neural network architectures, and the concept of training loops. This lesson will build upon that knowledge, focusing specifically on how to prepare your data so that PyTorch's data loading utilities can efficiently feed it to your models. This involves understanding the Dataset and DataLoader classes, which are fundamental to managing data pipelines in PyTorch. Imagine you're working with a collection of protein structures (e.g., PDB files) and their associated binding affinities for a particular ligand. Your goal might be to predict binding affinity from structural features. This type of data doesn't fit neatly into standard image or text datasets. You'll need to define how to load each protein structure, extract relevant features (e.g., atom coordinates, residue types, solvent accessibility), and pair them with their corresponding binding affinity labels. This is precisely where custom datasets come into play.
Building Custom Datasets with torch.utils.data.Dataset
The cornerstone of custom data handling in PyTorch is the torch.utils.data.Dataset abstract class. To create your own custom dataset, you'll inherit from this class and override two essential methods: __len__(self) : This method should return the total number of samples in your dataset. PyTorch uses this to determine the dataset's size. __getitem__(self, idx) : This method is responsible for loading and returning a single sample from your dataset at the given index idx . The sample typically consists of a data point (e.g., protein features) and its corresponding label (e.g., binding affinity). Let's consider a simplified example where we have a dataset of protein sequences and their stability scores. We'll assume the data is stored in a CSV file. import torch from torch.utils.data import Dataset, DataLoader import pandas as pd import numpy as np # A hypothetical function to convert a protein sequence to a numerical feature vector # In a real scenario, this would be a much more complex featurization process def featurize_sequence(sequence): # Example: one-hot encoding for simplicity, or embedding lookup amino_acids = 'ACDEFGHIKLMNPQRSTVWXY' feature_vector = [amino_acids.find(aa) for aa in sequence] # Pad or truncate to a fixed length if necessary max_len = 10 # Example fixed length if len(feature_vector) Once you have a custom Dataset , you'll typically wrap it with a torch.utils.data.DataLoader . The DataLoader handles batching, shuffling, and multi-process data loading, making your training process much more efficient. It iterates over your dataset, yielding batches of data that can be directly fed into your neural network. # Continue from the previous code block # Create a DataLoader batch_size = 2 data_loader = DataLoader(protein_dataset, batch_size=batch_size, shuffle=True) # Iterate through the DataLoader print("\nIterating through DataLoader:") for i, (batch_features, batch_labels) in enumerate(data_loader): print(f"Batch {i+1}:") print(f" Features shape: {batch_features.shape}") print(f" Labels shape: {batch_labels.shape}") # In a real scenario, you would pass batch_features to your model # and compare its output with batch_labels to calculate loss. The DataLoader is particularly important for large datasets that cannot fit entirely into memory. It loads data samples on demand, reducing memory footprint and speeding up training by potentially using multiple worker processes (controlled by the num_workers argument).
Key Takeaways:
Custom datasets are essential for real-world biological and drug discovery problems where standard datasets are insufficient. Inherit from torch.utils.data.Dataset to create your custom dataset. Implement __len__(self) to return the total number of samples. Implement __getitem__(self, idx) to load and return a single data-label pair as PyTorch tensors. Use torch.utils.data.DataLoader to efficiently batch, shuffle, and load data from your custom dataset during training. The DataLoader abstracts away the complexities of data iteration, making your training loop cleaner and more efficient.
Practice Exercise:
Imagine you have a collection of small molecule SMILES strings along with their predicted solubility values. Your task is to create a custom PyTorch Dataset for this data. Assume the data is in a CSV file named molecules.csv with columns 'SMILES' and 'Solubility'. For the purpose of this exercise, you can use a placeholder function for featurizing SMILES strings (e.g., converting each character to an integer ID, padding to a fixed length). Instantiate your dataset and then create a DataLoader with a batch size of 4, ensuring the data is shuffled. Print the shapes of the features and labels for the first batch.
Watch the full lesson — free
This topic is part of Structural Biology & Drug Discovery, a complete AI-narrated video course. Press play once and watch the entire lecture like a movie.
Start the course free →