Method test_lr_finder

tests/optimizers/test_lr_finder.py:65–104

(self)

        )

    def test_lr_finder(self):
        # 0.001 gives 54 examples
        with skip_if_downloading_fails():
            train_ds = MedNISTDataset(
                root_dir=self.root_dir,
                transform=self.transforms,
                section="validation",
                val_frac=0.001,
                download=True,
                num_workers=2,
            )
        train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=2)
        num_classes = train_ds.get_num_classes()

        model = DenseNet(
            spatial_dims=2, in_channels=1, out_channels=num_classes, init_features=2, growth_rate=2, block_config=(2,)
        )
        loss_function = torch.nn.CrossEntropyLoss()
        learning_rate = 1e-5
        optimizer = torch.optim.Adam(model.parameters(), learning_rate)

        lr_finder = LearningRateFinder(
            model=model,
            optimizer=optimizer,
            criterion=loss_function,
            device=device,
            pickle_module=pickle,
            pickle_protocol=4,
        )
        lr_finder.range_test(train_loader, val_loader=train_loader, end_lr=10.0, num_iter=5)
        print(lr_finder.get_steepest_gradient(0, 0)[0])

        if has_matplotlib:
            ax = plt.subplot()
            plt.show(block=False)
            lr_finder.plot(0, 0, ax=ax)  # to inspect the loss-learning rate graph
            plt.pause(3)
            plt.close()

        lr_finder.reset()  # to reset the model and optimizer to their initial state


if __name__ == "__main__":