\n",
+ "Task 1.2\n",
+ "Setup the data loader and log several batches to tensorboard.\n",
+ "
`\n",
+ "\n",
+ "VisCy builds on top of PyTorch Lightning. PyTorch Lightning is a thin wrapper around PyTorch that allows rapid experimentation. It provides a [DataModule](https://lightning.ai/docs/pytorch/stable/data/datamodule.html) to handle loading and processing of data during training. VisCy provides a child class, `HCSDataModule` to make it intuitve to access data stored in the HCS layout.\n",
+ " \n",
+ "The dataloader in `HCSDataModule` returns a batch of samples. A `batch` is a list of dictionaries. The length of the list is equal to the batch size. Each dictionary consists of following key-value pairs.\n",
+ "- `source`: the input image, a tensor of size 1*1*Y*X\n",
+ "- `target`: the target image, a tensor of size 2*1*Y*X\n",
+ "- `index` : the tuple of (location of field in HCS layout, time, and z-slice) of the sample."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "05ce4dbb",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Define a function to write a batch to tensorboard log.\n",
+ "\n",
+ "\n",
+ "def log_batch_tensorboard(batch, batchno, writer, card_name):\n",
+ " \"\"\"\n",
+ " Logs a batch of images to TensorBoard.\n",
+ "\n",
+ " Args:\n",
+ " batch (dict): A dictionary containing the batch of images to be logged.\n",
+ " writer (SummaryWriter): A TensorBoard SummaryWriter object.\n",
+ " card_name (str): The name of the card to be displayed in TensorBoard.\n",
+ "\n",
+ " Returns:\n",
+ " None\n",
+ " \"\"\"\n",
+ " batch_phase = batch[\"source\"][:, :, 0, :, :] # batch_size x z_size x Y x X tensor.\n",
+ " batch_membrane = batch[\"target\"][:, 1, 0, :, :].unsqueeze(\n",
+ " 1\n",
+ " ) # batch_size x 1 x Y x X tensor.\n",
+ " batch_nuclei = batch[\"target\"][:, 0, 0, :, :].unsqueeze(\n",
+ " 1\n",
+ " ) # batch_size x 1 x Y x X tensor.\n",
+ "\n",
+ " p1, p99 = np.percentile(batch_membrane, (0.1, 99.9))\n",
+ " batch_membrane = np.clip((batch_membrane - p1) / (p99 - p1), 0, 1)\n",
+ "\n",
+ " p1, p99 = np.percentile(batch_nuclei, (0.1, 99.9))\n",
+ " batch_nuclei = np.clip((batch_nuclei - p1) / (p99 - p1), 0, 1)\n",
+ "\n",
+ " p1, p99 = np.percentile(batch_phase, (0.1, 99.9))\n",
+ " batch_phase = np.clip((batch_phase - p1) / (p99 - p1), 0, 1)\n",
+ "\n",
+ " [N, C, H, W] = batch_phase.shape\n",
+ " interleaved_images = torch.zeros((3 * N, C, H, W), dtype=batch_phase.dtype)\n",
+ " interleaved_images[0::3, :] = batch_phase\n",
+ " interleaved_images[1::3, :] = batch_nuclei\n",
+ " interleaved_images[2::3, :] = batch_membrane\n",
+ "\n",
+ " grid = torchvision.utils.make_grid(interleaved_images, nrow=3)\n",
+ "\n",
+ " # add the grid to tensorboard\n",
+ " writer.add_image(card_name, grid, batchno)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f627e8e8",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "\n",
+ "# Initialize the data module.\n",
+ "\n",
+ "BATCH_SIZE = 42\n",
+ "# 42 is a perfectly reasonable batch size. After all, it is the answer to the ultimate question of life, the universe and everything.\n",
+ "# More seriously, batch size does not have to be a power of 2.\n",
+ "# See: https://sebastianraschka.com/blog/2022/batch-size-2.html\n",
+ "\n",
+ "data_module = HCSDataModule(\n",
+ " data_path,\n",
+ " source_channel=\"Phase\",\n",
+ " target_channel=[\"Nuclei\", \"Membrane\"],\n",
+ " z_window_size=1,\n",
+ " split_ratio=0.8,\n",
+ " batch_size=BATCH_SIZE,\n",
+ " num_workers=8,\n",
+ " architecture=\"2D\",\n",
+ " yx_patch_size=(512, 512), # larger patch size makes it easy to see augmentations.\n",
+ " augment=False, # Turn off augmentation for now.\n",
+ ")\n",
+ "data_module.setup(\"fit\")\n",
+ "\n",
+ "print(\n",
+ " f\"FOVs in training set: {len(data_module.train_dataset)}, FOVs in validation set:{len(data_module.val_dataset)}\"\n",
+ ")\n",
+ "train_dataloader = data_module.train_dataloader()\n",
+ "\n",
+ "# Instantiate the tensorboard SummaryWriter, logs the first batch and then iterates through all the batches and logs them to tensorboard.\n",
+ "\n",
+ "writer = SummaryWriter(log_dir=f\"{log_dir}/view_batch\")\n",
+ "# Draw a batch and write to tensorboard.\n",
+ "batch = next(iter(train_dataloader))\n",
+ "log_batch_tensorboard(batch, 0, writer, \"augmentation/none\")\n",
+ "\n",
+ "# Iterate through all the batches and log them to tensorboard.\n",
+ "for i, batch in enumerate(train_dataloader):\n",
+ " log_batch_tensorboard(batch, i, writer, \"augmentation/none\")\n",
+ "writer.close()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "f6ffc7d1",
+ "metadata": {
+ "cell_marker": "\"\"\""
+ },
+ "source": [
+ "There are multiple ways of seeing the tensorboard.\n",
+ "1. Jupyter lab forwards the tensorboard port to the browser. Go to http://localhost:6006/ to see the tensorboard.\n",
+ "2. You likely have an open viewer in the first cell where you loaded tensorboard jupyter extension.\n",
+ "3. If you want to see tensorboard in a specific cell, use the following code.\n",
+ "```\n",
+ "notebook.list() # View open TensorBoard instances\n",
+ "notebook.display(port=6006, height=800) # Display the TensorBoard instance specified by the port.\n",
+ "```"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ed1448ff",
+ "metadata": {
+ "cell_marker": "\"\"\"",
+ "lines_to_next_cell": 0
+ },
+ "source": [
+ "## View augmentations using tensorboard.\n",
+ "\n",
+ "\n",
+ "Task 1.3\n",
+ "Turn on augmentation and view the batch in tensorboard."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "245ae4dd",
+ "metadata": {
+ "lines_to_next_cell": 2
+ },
+ "outputs": [],
+ "source": [
+ "##########################\n",
+ "######## TODO ########\n",
+ "##########################\n",
+ "\n",
+ "# Write code to turn on augmentations, change batch sizes and log them to tensorboard.\n",
+ "# See how the training data changes as a function of these parameters.\n",
+ "# Remember to call `data_module.setup(\"fit\")` after changing the parameters."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "e14f4faf",
+ "metadata": {
+ "tags": [
+ "solution"
+ ]
+ },
+ "outputs": [],
+ "source": [
+ "##########################\n",
+ "######## Solution ########\n",
+ "##########################\n",
+ "\n",
+ "data_module.augment = True\n",
+ "data_module.batch_size = 21\n",
+ "data_module.split_ratio = 0.8\n",
+ "data_module.setup(\"fit\")\n",
+ "\n",
+ "train_dataloader = data_module.train_dataloader()\n",
+ "# Draw batches and write to tensorboard\n",
+ "writer = SummaryWriter(log_dir=f\"{log_dir}/view_batch\")\n",
+ "for i, batch in enumerate(train_dataloader):\n",
+ " log_batch_tensorboard(batch, i, writer, \"augmentation/some\")\n",
+ "writer.close()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "85743be9",
+ "metadata": {
+ "cell_marker": "\"\"\""
+ },
+ "source": [
+ "## Construct a 2D U-Net for image translation.\n",
+ "See ``viscy.unet.networks.Unet2D.Unet2d`` for configuration details.\n",
+ "We setup a fresh data module and instantiate the trainer class."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "44fddf02",
+ "metadata": {
+ "lines_to_next_cell": 2
+ },
+ "outputs": [],
+ "source": [
+ "\n",
+ "# The entire training loop is contained in this cell.\n",
+ "\n",
+ "GPU_ID = 0\n",
+ "BATCH_SIZE = 10\n",
+ "YX_PATCH_SIZE = (512, 512)\n",
+ "\n",
+ "\n",
+ "# Dictionary that specifies key parameters of the model.\n",
+ "phase2fluor_config = {\n",
+ " \"architecture\": \"2D\",\n",
+ " \"num_filters\": [24, 48, 96, 192, 384],\n",
+ " \"in_channels\": 1,\n",
+ " \"out_channels\": 2,\n",
+ " \"residual\": True,\n",
+ " \"dropout\": 0.1, # dropout randomly turns off weights to avoid overfitting of the model to data.\n",
+ " \"task\": \"reg\", # reg = regression task.\n",
+ "}\n",
+ "\n",
+ "phase2fluor_model = VSUNet(\n",
+ " model_config=phase2fluor_config.copy(),\n",
+ " batch_size=BATCH_SIZE,\n",
+ " loss_function=torch.nn.functional.l1_loss,\n",
+ " schedule=\"WarmupCosine\",\n",
+ " log_num_samples=10, # Number of samples from each batch to log to tensorboard.\n",
+ " example_input_yx_shape=YX_PATCH_SIZE,\n",
+ ")\n",
+ "\n",
+ "# Reinitialize the data module.\n",
+ "phase2fluor_data = HCSDataModule(\n",
+ " data_path,\n",
+ " source_channel=\"Phase\",\n",
+ " target_channel=[\"Nuclei\", \"Membrane\"],\n",
+ " z_window_size=1,\n",
+ " split_ratio=0.8,\n",
+ " batch_size=BATCH_SIZE,\n",
+ " num_workers=8,\n",
+ " architecture=\"2D\",\n",
+ " yx_patch_size=YX_PATCH_SIZE,\n",
+ " augment=True,\n",
+ ")\n",
+ "phase2fluor_data.setup(\"fit\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "292cb22a",
+ "metadata": {
+ "cell_marker": "\"\"\""
+ },
+ "source": [
+ "
\n",
+ "Task 1.4\n",
+ "Setup the training for ~30 epochs\n",
+ "
\n",
+ "\n",
+ "Tips:\n",
+ "- Set ``default_root_dir`` to store the logs and checkpoints\n",
+ "in a specific directory."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b9cf215b",
+ "metadata": {
+ "lines_to_next_cell": 2,
+ "title": "Setup trainer and check for errors."
+ },
+ "outputs": [],
+ "source": [
+ "\n",
+ "# fast_dev_run runs a single batch of data through the model to check for errors.\n",
+ "trainer = VSTrainer(accelerator=\"gpu\", devices=[GPU_ID], fast_dev_run=True)\n",
+ "\n",
+ "# trainer class takes the model and the data module as inputs.\n",
+ "trainer.fit(phase2fluor_model, datamodule=phase2fluor_data)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "5f11d604",
+ "metadata": {
+ "lines_to_next_cell": 2
+ },
+ "outputs": [],
+ "source": [
+ "\n",
+ "GPU_ID = 0\n",
+ "n_samples = len(phase2fluor_data.train_dataset)\n",
+ "steps_per_epoch = n_samples // BATCH_SIZE # steps per epoch.\n",
+ "n_epochs = 30\n",
+ "\n",
+ "trainer = VSTrainer(\n",
+ " accelerator=\"gpu\",\n",
+ " devices=[GPU_ID],\n",
+ " max_epochs=n_epochs,\n",
+ " log_every_n_steps=steps_per_epoch // 2,\n",
+ " # log losses and image samples 2 times per epoch.\n",
+ " default_root_dir=Path(\n",
+ " log_dir, \"phase2fluor\"\n",
+ " ), # lightning trainer transparently saves logs and model checkpoints in this directory.\n",
+ ")\n",
+ "\n",
+ "# Log graph\n",
+ "trainer.logger.log_graph(phase2fluor_model, phase2fluor_data.train_dataset[0][\"source\"])\n",
+ "# Launch training.\n",
+ "trainer.fit(phase2fluor_model, datamodule=phase2fluor_data)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "157a989c",
+ "metadata": {
+ "cell_marker": "\"\"\""
+ },
+ "source": [
+ "
\n",
+ "Checkpoint 1\n",
+ "\n",
+ "Now the training has started,\n",
+ "we can come back after a while and evaluate the performance!\n",
+ "
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "d68dab93",
+ "metadata": {
+ "cell_marker": "\"\"\"",
+ "incorrectly_encoded_metadata": "id='1_fluor2phase'>",
+ "title": "
\n",
+ "Checkpoint 2\n",
+ "Please summarize hyperparameters and performance of your models in [this google doc](https://docs.google.com/document/d/1hZWSVRvt9KJEdYu7ib-vFBqAVQRYL8cWaP_vFznu7D8/edit#heading=h.n5u485pmzv2z)\n",
+ "\n",
+ "Now that you have trained two models, let's think about the following questions:\n",
+ "- What is the information content of each channel in the dataset?\n",
+ "- How would you use image translation models?\n",
+ "- What can you try to improve the performance of each model?\n",
+ "\n",
+ "\n",
+ "