Elegant Configurations in Python

I’ve always struggled to manage all the hyperparameters in machine learning projects. Over the years, I have tried many different options including YAML files, Python dictionaries, dataclasses, argparse. Each approach fell short in one way or another.

Maybe I’m a bit picky, but when it comes to hyperparameters I want

• A single, centralized place to define them
• The ability to easily override individual values from the command line (sorry YAML and dataclasses)
• Auto-complete and typo protection (sorry dictionaries and argparse)

Of course there are massive hard-to-understand codebases out there that attempt to solve this problem. But I’m not looking for complexity, I want something small, elegant, and easy to understand. Our goal for this blog post is to build exactly that.

(Spoiler: here it is.)

Defining the Configuration

Our first goal is to define all our hyperparameters as well as default values. For simple settings we can do that easily using a class. In particular, a dataclass allows us to do this with minimal duplication and boilerplate:

@dataclass
class Configuration:
    dataset_name: str = "CIFAR-10"
    epochs: int = 24
    use_test_set: bool = False
    ...

Using this class allows our IDE to generate auto-completion and warn us about typos in our code. We can also use the class for type-hinting.

Adding Command Line Parsing

In Python, argparse is the standard module for parsing command line arguments. To keep our code simple, we don’t want to reinvent the wheel, so we will use this package. There is nothing most coders hate more than duplication, so we will try to avoid it. So we’ll generate our parser using the metadata provided by the dataclass.

Luckily, the dataclass itself provides the required information such as the field types. Using this, we can define our parser in the following way:

def create_parser(cls) -> argparse.ArgumentParser:
    if not dataclasses.is_dataclass(cls):
        raise TypeError(f"{cls.__name__} must be a dataclass")
        
    parser = argparse.ArgumentParser()
    for field in dataclasses.fields(cls):
        arg_name = f"--{field.name}".replace("_", "-")
        arg_type = field.type
        arg_default = field.default
				
        if arg_type == bool:
            action = "store_true" if arg_default is False else "store_false"
            parser.add_argument(*arg_name, action=action)
        else:
            parser.add_argument(*arg_name, type=arg_type, default=arg_default)
    return parser    

Now we can parse our command line inputs. Goal achieved. 🎉 However, parser.parse_args() gives us an argparse.Namespace object, but we want an instance of the Configuration class. In order to get that we define the function below.

def from_command_line(cls):
    parser = create_parser(cls)
    args = parser.parse_args()
    return cls(**vars(args))

In our code we also define a similar class method from_command_line for our configuration. The class method has the advantage that PyCharm infers the correct type.

How to use the Configuration:

The code above accomplishes our goal. In a few lines of code, we used the power of dataclasses and argparse to create an elegant way of defining and changing the hyperparameters for our scripts.

We can obtain an instance of the configuration class easily either using command line arguments, or by explicitly instantiating it using the class constructor. This also makes it convenient to use in unit tests.

config1 = Configuration.from_command_line()
config2 = Configuration(learning_rate = 3e-4)
config3 = parsing.from_command_line(Configuration)

Furthermore, we can also use our class for type hinting:

def train(c: Configuration):
    ...

Extra 1: Short flags.

Another thing most programmers dread is redundant typing (the kind that is done on a keyboard). We can free ourselves from a bit of typing by creating short flags for the arguments. One way is as follows: We first define a decorator, that allows us to define and save some short flags for our Configuration class. We also slightly adapt the create_parser function to include those short flags.

def with_short_flags(**kwargs):
    def decorator(cls):
        setattr(cls, "__short_flags__", kwargs)
        return cls

    return decorator


def create_parser(cls) -> argparse.ArgumentParser:
    if not dataclasses.is_dataclass(cls):
        raise TypeError(f"{cls.__name__} must be a dataclass")

    parser = argparse.ArgumentParser()
    short_flags = getattr(cls, "__short_flags__", {})

    for field in dataclasses.fields(cls):
        arg_name = f"--{field.name}".replace("_", "-")
        arg_type = field.type
        arg_default = field.default

        if field.name in short_flags:
            flags = [short_flags[field.name], arg_name]
        else:
            flags = [arg_name]

        if arg_type == bool:
            action = {True: "store_false", False: "store_true"}[arg_default]
            parser.add_argument(*flags, action=action)
        else:
            parser.add_argument(*flags, type=arg_type, default=arg_default)
    return parser

With this, we can define our configuration in the following way:

@parsing.with_short_flags(dataset_name="-d", epochs="-e")
@dataclass
class Configuration:
    dataset_name: str = "CIFAR-10"
    epochs: int = 24
    use_test_set: bool = False

This full version is of course also available on Github.

Conclusion & Lesson Learned

I struggled for some time to create a simple, elegant way of defining hyperparameters that allows autocompletion and setting values the command line. I think the solution described above is finally something I am happy with - at least for now. For small projects with only a few hyperparameters it seems to me to strike a good balance between simplicity and power. (For large project, however, it is probably still better to use and existing framework such as Hydra.)

For our solution, we used what Python already provides: Dataclasses offer readable definitions and autocomplete. The argparse module already provides us with argument parsing, including conveniences such as helpful error messages. Our job was just to connect these two established components. This approach worked out much better than many of my previous attempts, where I tried to do everything from scratch, reinventing the wheel multiple times.

Finally, if needed we could also easily extend our solution with a few adaptations to work for nested configurations or load YAML files. Also, have a look at some of the alternatives.

I am happy to hear your thoughts!