Repository: oumi-ai/oumi
File: tests/integration/models/test_cnn_classifier.py (Python)

Code before:

        training_params = TrainingParams(
            output_dir=str(Path(tmp_dir) / "mnist_cnn_test_output"),
            per_device_train_batch_size=32,
            per_device_eval_batch_size=8,
            eval_strategy="steps",
            eval_steps=50,
            max_steps=100,
            save_steps=0,
            logging_steps=50,
            optimizer="adam",
            learning_rate=2e-3,
        )

        trainer = Trainer(
            model=model,
            tokenizer=None,  # No tokenizer! The custom model is non-textual
            args=training_params,
            train_dataset=train_dataset,
            dataloader_num_workers=2,
            dataloader_prefetch_factor=32,
        )

        trainer.train()

        with torch.no_grad():
            model_device = next(model.parameters()).device
            batch_size = 2
            for test_input in DataLoader(test_dataset, batch_size=batch_size):
                inputs = _convert_example_to_model_input(
                    test_input, device=model_device
                )

                # Without label
                outputs = trainer.model(images=inputs["images"])
                assert outputs.keys() == ({"logits"})
                assert isinstance(outputs["logits"], torch.Tensor)
                logits = outputs["logits"].cpu().numpy()
                assert logits.shape == (batch_size, 10)
                assert logits.dtype == np.float32

                # With label
                outputs = trainer.model(
                    images=inputs["images"], labels=inputs["labels"]
                )
                assert outputs.keys() == ({"logits", "loss"})
                assert isinstance(outputs["logits"], torch.Tensor)
                logits = outputs["logits"].cpu().numpy()
                assert logits.shape == (batch_size, 10)
                assert logits.dtype == np.float32
                loss = outputs["loss"].cpu().numpy()

Review comment (bug):
nit: we could split into multiple separate tests for easier debugging. This tests seems to test a lot of

Code after:

def _validate_prediction(model: torch.nn.Module, dataset, *, max_iters: int = 5):
    with torch.no_grad():
        model_device = next(model.parameters()).device
        batch_size = 2

        for idx, test_input in enumerate(DataLoader(dataset, batch_size=batch_size)):
            if idx >= max_iters:
                break
            inputs = _convert_example_to_model_input(test_input, device=model_device)

            # Without label
            outputs = model(images=inputs["images"])
            assert outputs.keys() == ({"logits"})
            assert isinstance(outputs["logits"], torch.Tensor)
            logits = outputs["logits"].cpu().numpy()
            assert logits.shape == (batch_size, 10)
            assert logits.dtype == np.float32

            # With label
            outputs = model(images=inputs["images"], labels=inputs["labels"])
            assert outputs.keys() == ({"logits", "loss"})
            assert isinstance(outputs["logits"], torch.Tensor)
            logits = outputs["logits"].cpu().numpy()
            assert logits.shape == (batch_size, 10)
            assert logits.dtype == np.float32
            loss = outputs["loss"].cpu().numpy()
            assert loss.shape == ()
            assert loss.dtype == np.float32


def test_basic_training_and_prediction():
    test_data = _create_test_data()

    # Create a model.
    model_params = ModelParams(
        model_name="CnnClassifier",
        load_pretrained_weights=False,
        model_kwargs={
            "image_width": 28,
            "image_height": 28,
            "in_channels": 1,
            "output_dim": 10,
        },
    )
    model = build_model(model_params)

    with (
        tempfile.TemporaryDirectory() as tmp_dir,
    ):
        training_params = TrainingParams(