Spaces:
Sleeping
Sleeping
File size: 9,490 Bytes
1f0bc98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 |
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Building a PyTorch Training Loop"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In order to be able to access the data on Hugging Face Hub and build the\n",
"data loaders for our training loop, we should import the necessary libraries\n",
"first"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from datasets import load_dataset # Loading datasets from Hugging Face Hub\n",
"import torch # PyTorch\n",
"from torch.utils.data import DataLoader # PyTorch DataLoader for creating batches\n",
"from pprint import pprint # Pretty print\n",
"from tqdm import tqdm # Progress bar"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this tutorial, we are going to work with the\n",
"[PubChemQC-B3LYP/6-31G*//PM6\n",
"Dataset](https://huggingface.co/datasets/molssiai-hub/pubchemqc-b3lyp)\n",
"(PubChemQC-B3LYP for short) from the [PubChemQC dataset\n",
"collection](https://huggingface.co/collections/molssiai-hub/pubchemqc-datasets-669e5482260861ba7cce3d1c).\n",
"Let us set a few variables and load the dataset as shown below"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After importing the modules, we set a few variables that will be used throughout\n",
"this demo."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# path to the dataset repository on the Hugging Face Hub\n",
"path = \"molssiai-hub/pubchemqc-b3lyp\"\n",
"\n",
"# set the dataset configuration/subset name\n",
"name = \"b3lyp_pm6\"\n",
"\n",
"# set the dataset split\n",
"split = \"train\"\n",
"\n",
"# load the dataset\n",
"hub_dataset = load_dataset(path=path,\n",
" name=name,\n",
" split=split,\n",
" streaming=True,\n",
" trust_remote_code=True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Here, we set the `streaming` parameter to `True` to avoid downloading the\n",
"dataset on disk and ensure streaming the data from the hub. In this mode, the\n",
"`load_dataset` function returns an `IterableDataset` object that can be iterated\n",
"over and provide access to the data. The `trust_remote_code` argument is also\n",
"set to `True` to allow the usage of a custom [load\n",
"script](https://huggingface.co/datasets/molssiai-hub/pubchemqc-b3lyp/blob/main/pubchemqc-b3lyp.py)\n",
"for the data."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"By default, the Hugging Face data objects' `__getitem__` method returns a native\n",
"Python object (e.g., a dictionary). However, we can use the `with_format()`\n",
"method to specify the format of the data we want to access. In our case, we want\n",
"to use the `torch.tensor` format to build the data loaders for our training\n",
"loop. Let us transform our data and check the result"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set the dataset format to PyTorch tensors\n",
"hub_dataset = hub_dataset.with_format(\"torch\")\n",
"\n",
"# fetch the first data point\n",
"next(iter(hub_dataset.take(1)))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see that the type of the numerical features in our data sample are\n",
"transformed to `torch.tensor` objects. Let us access the `coordinates` field\n",
"to make this more clear"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# fetch the first data point\n",
"data_point = next(iter(hub_dataset.take(1)))\n",
"\n",
"# print the coordinates of the first data point and its type\n",
"data_point[\"coordinates\"], type(data_point[\"coordinates\"])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In the code snippet above, we have wrapped the `IterableDataset` object, `hub_dataset`,\n",
"inside an `iter()` function to create an iterator object and used the `next()` function\n",
"to iterate once over it and access the first sample in it."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Our PubChemQC-B3LYP `IterableDataset` object is divided into multiple shards\n",
"to enable multiprocessing and help shuffling the data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(f\"the PubChemQC-B3LYP dataset has {hub_dataset.n_shards} shards\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If we want to shuffle our data, the shards will also be shuffled. This is\n",
"important to consider when building the PyTorch data loaders for our training\n",
"loop."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# shuffle the dataset\n",
"hub_dataset = hub_dataset.shuffle(seed=123, buffer_size=1000)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `buffer_size` controls the size of a container object from which we randomly\n",
"sample examples from. For instance, when we call the `IterableDataset.shuffle()`\n",
"function, the first thousand examples in the buffer are randomly sampled and the\n",
"selected examples in the buffer are then replaced with new examples from the\n",
"dataset. The `buffer_size` argument is set to 1000 by default. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"A nice feature of the Hugging Face dataset objects is that they can be directly\n",
"passed to PyTorch DataLoaders as shown below"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# create a PyTorch DataLoader with a batch size of 4\n",
"dataloader = DataLoader(hub_dataset, batch_size=4, collate_fn=lambda x: x)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"By default, the `DataLoader` object will use a default collator function which\n",
"creates batches of data and transforms them into `torch.tensors`. For our\n",
"dataset examples, however, we cannot use the default collator function because\n",
"our data samples are not of the same length (different molecules may have\n",
"different number of atoms and coordinates). To circumvent this problem, we can\n",
"define a lambda function that yields each data point, which is a dictionary,\n",
"without any transformation."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Similar to the `hub_dataset`, we can also wrap the `dataloader` object inside an\n",
"iterator and use the `next()` function to access the first batch of data "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_point = next(iter(dataloader))\n",
"data_point[0][\"coordinates\"]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Building a Training Loop in PyTorch"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we know how to access, fetch and shuffle batches of data samples in our\n",
"PyTorch data loader, we can build a simple training loop to train a model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# set up the training loop\n",
"for epoch in range(1, 4, 1):\n",
"\n",
" # set the epoch\n",
" hub_dataset.set_epoch(epoch)\n",
"\n",
" # iterate over the batches in the DataLoader\n",
" for i, batch in enumerate(tqdm(dataloader, total=4, desc=f\"Epoch {epoch}\")):\n",
" if i == 4:\n",
" pprint(f\"The isomeric SMILES from the first data point of the {i}th batch: {batch[0]['pubchem-isomeric-smiles']}\",\n",
" width=100,\n",
" compact=True)\n",
" break\n",
" print(f\"Epoch: {epoch}, Batch: {i+1}, Batch size: {len(batch)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In the code snippet above, we have used the `set_epoch(epoch)` function which\n",
"is often used with PyTorch data loaders and in distributed settings to augment the\n",
"random seed for reshuffling at the beginning of each epoch."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "hugface",
"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.10.13"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|