import unittest
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler
import importlib
import io
import sys
import pandas as pd


def _dynamic(test, cond, success, fail):
    if cond:
        test._testMethodName = success
        test.assertTrue(True)
    else:
        test._testMethodName = fail
        test.fail(fail)


class TestUserCode(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.df = pd.read_csv(
            "https://content-media-cdn.codefinity.com/courses/1dd2b0f6-6ec0-40e6-a570-ed0ac2209666/section_2/iris.csv"
        )
        cls.features = torch.tensor(cls.df.drop(columns="species").values, dtype=torch.float32)
        cls.targets = torch.tensor(cls.df["species"].values, dtype=torch.long)

    # -------- Dataset exists --------
    def test_dataset_exists(self):
        import user_code
        cond = hasattr(user_code, "iris_dataset")
        _dynamic(
            self,
            cond,
            "`iris_dataset` is successfully created.",
            "`iris_dataset` is missing. Create it using `TensorDataset(features_tensor, target_tensor)`."
        )

    # -------- Dataset is TensorDataset --------
    def test_dataset_is_tensor_dataset(self):
        import user_code
        cond = isinstance(user_code.iris_dataset, TensorDataset)
        _dynamic(
            self,
            cond,
            "`iris_dataset` is a valid TensorDataset object.",
            "`iris_dataset` must be created using `TensorDataset(features_tensor, target_tensor)`."
        )

    # -------- Dataset content correctness --------
    def test_dataset_tensors_correct(self):
        import user_code
        ds = user_code.iris_dataset

        try:
            feats_ok = torch.allclose(ds.tensors[0], self.features)
            targets_ok = torch.allclose(ds.tensors[1], self.targets)
            cond = feats_ok and targets_ok
        except Exception:
            cond = False

        _dynamic(
            self,
            cond,
            "`iris_dataset` correctly stores the feature and target tensors.",
            "`iris_dataset` must contain `(features_tensor, target_tensor)` in that order."
        )

    # -------- DataLoader exists --------
    def test_dataloader_exists(self):
        import user_code
        cond = hasattr(user_code, "iris_loader")
        _dynamic(
            self,
            cond,
            "`iris_loader` is successfully created.",
            "`iris_loader` is missing. Wrap your dataset using `DataLoader(...)`."
        )

    # -------- DataLoader parameters --------
    def test_dataloader_params(self):
        import user_code
        loader = user_code.iris_loader

        cond = (
            isinstance(loader, DataLoader)
            and loader.batch_size == 32
            and isinstance(loader.sampler, RandomSampler)
        )

        _dynamic(
            self,
            cond,
            "`iris_loader` uses batch_size=32 and shuffles the data (RandomSampler).",
            "`iris_loader` must be created with `batch_size=32` and `shuffle=True` so that it uses a RandomSampler."
        )

    # -------- First batch correctness --------
    def test_first_batch_shape(self):
        import user_code

        loader = user_code.iris_loader
        batch_features, batch_targets = next(iter(loader))

        cond = (
            batch_features.ndim == 2
            and batch_features.shape[1] == 4
            and batch_targets.ndim == 1
        )

        _dynamic(
            self,
            cond,
            "Batches from `iris_loader` have the expected shapes.",
            "Batch shapes are incorrect. Features must be `[batch_size, 4]` and targets must be `[batch_size]`."
        )

    # -------- Output printing --------
    def test_output_printed(self):
        import user_code

        buf = io.StringIO()
        sys.stdout = buf
        importlib.reload(user_code)
        sys.stdout = sys.__stdout__

        cond = len(buf.getvalue().strip()) > 0

        _dynamic(
            self,
            cond,
            "The DataLoader output is printed successfully.",
            "No output detected. Ensure you print the first batch from the DataLoader."
        )


if __name__ == "__main__":
    unittest.main()


test_code.py

Learn the fundamental and advanced concepts needed to work with PyTorch efficiently. Gain a solid understanding tensors, including creation, operations, and reshaping. Explore the essentials of gradients, backpropagation, and linear regression before moving on to handling datasets. Master the skill needed build, train, and evaluate neural networks.

Challenge: Preparing the Iris Dataset

Solution