Deep Learning with JAX and Flax/

...

/

Solution: Distributed Training with JAX and Flax

Solution: Distributed Training with JAX and Flax

Let's review the solution.

We'll cover the following...

Solution 1: Load the dataset

First, we define the dataset class like we did in previous chapters. However, this time for the Cars and Bikes dataset:

Python 3.8
from PIL import Image
import pandas as pd
from torch.utils.data import Dataset
class CarBikeDataset(Dataset):
  def __init__(self, root_dir, annotation_file, transform=None):
    self.root_dir = root_dir
    self.annotations = pd.read_csv(annotation_file)
    self.transform = transform
  def __len__(self):
    return len(self.annotations)
  def __getitem__(self, index):
    img_id = self.annotations.iloc[index, 0]
    img = Image.open(os.path.join(self.root_dir, img_id)).convert("RGB")
    y_label = torch.tensor(float(self.annotations.iloc[index, 1]))
    if self.transform is not None:
      img = self.transform(img)
    return (img, y_label)

In the code above:

  • Lines 6–9: The parametrized constructor defines the root directory containing the images, image labels (using the annotation file), and image transformations (transforms).
  • Line 11: The __len__ function returns the number of samples in our dataset.
  • Lines 14–22: The __getitem__ function finds the image in the dataset at index idx, converts it to an RGB, applies transformations (if any), and returns it along with its corresponding label that is converted to a tensor.

Visualizing the dataset

Let’s visualize one image from each class of the Cars and Bikes dataset.

Python 3.8
from PIL import Image
img1 = Image.open("/usercode/cat.1.png")
img1.save("output/img1.png")
img2 = Image.open("/usercode/dog.3.png")
img2.save("output/img2.png")

Next, we create a pandas DataFrame that will contain the categories.

Python 3.8
import os
import pandas as pd
train_df = pd.DataFrame(columns=["img_path","label"])
train_df["img_path"] = os.listdir("cars_and_bikes/")
for idx, i in enumerate(os.listdir("cars_and_bikes/")):
  if "car" in i:
    train_df["label"][idx] = 0
  if "bike" in i:
    train_df["label"][idx] = 1
train_df.to_csv (r'train_csv.csv', index = False, header=True)

In the code above:

  • Line 4: We initialize a DataFrame for storing paths and labels of our images.
  • Lines 5–11: We populate the DataFrame by getting the image paths from the directory
...