Enums parsing in hparams.yaml generated

[grajat90]

What does this PR do?

Fixes #8912

Does your PR introduce any breaking changes? If yes, please list them.

NONE

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Yesss

Make sure you had fun coding 🙃

[grajat90]

Passed functionality testing as well as unit testing when checked locally.
Does not break anything else when tested locally. Tested using following piece of code to run a test classifier on MNIST and have tensorboard log the hyperparams in the hparams.yaml file. This test fails on master and passes on this branch:

import pytorch_lightning as pl
from enum import Enum
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torch.utils.data import random_split
from torchvision.datasets import MNIST
from torchvision import transforms
from pytorch_lightning.loggers import TensorBoardLogger

logger = TensorBoardLogger(save_dir="tb_logs", name="my_model")

class Options(str, Enum):
    option1 = "option1"
    option2 = "option2"
    option3 = "option3"

class LitAutoEncoder(pl.LightningModule):
    def __init__(self,learning_rate: float = 0.0001,
                 switch: Options = Options.option3, # argument of interest
                 batch_size: int = 32):
        super().__init__()
        self.save_hyperparameters()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 64),
            nn.ReLU(),
            nn.Linear(64, 3))
        self.decoder = nn.Sequential(
            nn.Linear(3, 64),
            nn.ReLU(),
            nn.Linear(64, 28 * 28))
      

    def forward(self, x):
        embedding = self.encoder(x)
        return embedding

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)    
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y = val_batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('val_loss', loss)

# data
dataset = MNIST('', train=True, download=True, transform=transforms.ToTensor())
mnist_train, mnist_val = random_split(dataset, [55000, 5000])

train_loader = DataLoader(mnist_train, batch_size=32)
val_loader = DataLoader(mnist_val, batch_size=32)

# model
model = LitAutoEncoder()

# training
trainer = pl.Trainer(limit_train_batches=0.2, logger=logger, max_epochs=1)
trainer.fit(model, train_loader, val_loader)

[dasayan05]

Hi, I opened the issue #8912 and I would like to point out few things @tchaton

1. Please note that OmegaConf.save(..) writes the .name of an Enum variable, whereas this PR is trying to write the .value. This will cause inconsistency in the written yaml file.
2. Now if you follow OmegaConf's way of writing .name, then the tests here will not pass because the == operator tries to match .value of an enum.
3. I just noticed that the save_hparams_to_yaml(..) do not have a use_omegaconf=.. guard. So everytime the tests here are writing the yaml using (unless it is not installed at all) OmegaConf and reading without it.

[awaelchli]

@grajat90 let's add a changelog entry too :)

[grajat90]

@grajat90 let's add a changelog entry too :)

Hey sure

Also wanted some clarification:
The way omegaconf saves enum type hyperparams is it converts the chosen out of the enums to a string and just saves it as string. I’ve replicated the functionality. But I think it would be better if it were converted to a dict as you pointed out. What I’ve done locally to try out is created custom parsing for enums. I save the enum definition and recreate the enum when we load it back from the file, select the option from the enum and return it to the hparams object. Is this the approach I should pursue or just try to replicate how omegaconf does it?

[grajat90]

As part of this latest commit, I have created custom parsing for enum type objects to store them in the hparams.yaml file as a dict. This is read back, the entire enum is recreated according to the definition, and is returned in the new object, thus retaining all enum based functionality that might be occurring in the model (using dot operator or any other enum special functions). This appears to be the best way to tackle it, please advice.

[codecov]

Codecov Report

Merging #9170 (c83b9e2) into master (d9dfb2e) will decrease coverage by 4%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #9170    +/-   ##
=======================================
- Coverage      93%     89%    -4%     
=======================================
  Files         182     182            
  Lines       16140   16142     +2     
=======================================
- Hits        14930   14316   -614     
- Misses       1210    1826   +616     

there were substantial changes since

[grajat90]

Hey if anyone can let me know if this is a better approach or the earlier approach, I can just drop the commit and we can continue with the previous method.

[awaelchli]

@grajat90 I do not know which one is better, I'm not familiar with this enum logic, but the current approach does not look right to me due to access to private member attributes.

[tchaton]

LGTM !

[rohitgr7]

LG! just some minor doc updates.