Have you ever named your experiment <code>rnet50_sgd-cos-lr1e-2_imnet200_augs-ver8</code>?
Have you ever experienced that awkward silence when your manager asks you to reproduce an experiment,
which you cannot do, because you have overwritten original code and lost the initial value of a learning rate?
If your answer is yes, we have something for you.
<h1 id="best-practices">Best practices</h1>
In AI research, your efficiency depends on how quickly you implement and deploy ideas.
The efficiency of your team depends on how quickly they can reuse or reproduce things you crafted.
For that reason, your team should implement a standardized protocol of communication and exchange of experimental setups.
At LeapMind, we use hydra along with PyTorch Lightning templates which we believe is the best practice for organizing AI research.
<img src="/images/posts/techblog/20200828-joniak-hydra-lightning/hydra-lightning.png" alt="An overview of the environment we use at LeapMind.">
An overview of the environment we use at LeapMind.
<h1 id="organize-your-code-lightning">Organize your code: Lightning</h1>
Do you know this mantra?
Data, gpu, feed forward, loss, zero grad, backward and logging?
A lot of logging. Well, that’s the training loop which you saw a million times and which looks the same for every project.
With an exception that sometimes it might get enormously long and nasty when adding more and more functionalities.
Here is where Pytorch Lightning comes in.
Simply saying, it’s an organized way to write a training loop.
Saying precisely: it’s a way to separate research and engineering by decoupling data from model and model from the training mechanism.
Let’s look at an example.
<code>LightningModule</code>, the core of the Lightning, defines functions as: <code>configure_optimizers()</code>, <code>forward()</code>, <code>training_step()</code>, <code>training_epoch_end()</code>.
Does it sound familiar?
Yes, this is where you write logic you would normally put in the loop.
This time, however, everything is neatly organized and – hold on to your code – Lightning will launch and manage the training loop for you.
The <code>LightningModule</code> defines <code>forward()</code> – does it mean the definition of your <code>SuperResNet50</code> also goes there?
No! It stays where it’s always been, in <code>superesrnet50.py</code>, you only inject its dependency into the module and call it in the forward.
Similar happens to a dataset: it gets injected to a so called <code>DataModule</code>,
where only five methods are to be implemented: <code>setup()</code>, <code>prepare_data()</code>, <code>{train, val, test}_dataloader()</code>.
As one might guess, the Lightning engine will call these methods whenever needed – so you don’t need to worry about the training flow.
Do you see an advantage of the Lightning now?
Thanks to it, a model, data and training do not depend on each other, meaning you can easily replace any of them.
Just imagine how much flexibility it gives when trying out new network architectures!
We won’t dive into technical details on the Lighting as its official documentation is a great resource.
<h1 id="organize-your-experimental-setups-hydra">Organize your experimental setups: hydra</h1>
The most common way to pass parameters to your system is, unfortunately, the <code>argparse</code>.
Let’s face the truth: these projects are just nasty – <code>argparse</code>-ing takes a lot of code which is often totally unreadable.
On the top, how does one know what arguments to pass to reproduce your experiments? Wrap them up in a <code>.sh</code> file?
Do we add more <code>argparse</code>s as the project grows?
Certainly not.
Fortunately we have hydra.
Thanks to hydra we can elegantly separate experiment configuration from code – in the same way we decoupled model and data using lightning.
An example would tell it all.
Suppose we have a custom lightning module with following signature:
<pre><code class="hljs language-python">class SueperResNet50Module(pl.LightningModule):
def __init__(self, network: nn.Module, loss_fn: nn.Module, lr: float, epochs: int):
...
</code></pre>
Note it’s already the Lightning module, meaning we would define training logic in there.
To instantiate the module, we could write the following <code>.yaml</code> file:
<pre><code class="hljs language-yaml">_target_: src.pl_datamodules.imagenet.CustomImageNetDM

network:
	_target_: src.architecture.SuperResNet50
loss_fn:
	_target_: torch.nn.CrossEntropyLoss
lr: 1e-2
epoch: 100
</code></pre>
First, notice the special keyword, <code>_target_</code> which specifies instantiatable objects.
As a value, it takes a package name followed by a class name.
It can be either a library (eg. <code>torch.nn.CrossEntropyLoss</code>), or our own package (<code>src.pl_datamodules.imagenet.CustomImageNetDM</code>).
Second, hydra will automatically attempt to match targets with object argument names, i.e. config value <code>lr</code> will be passed to the <code>lr</code> argument, and so on.
Do you need to run an experiment with a different learning rate now?
It’s as easy as writing a new config file – clean and easy to reproduce.
But hold on, isn't it burdensome to create a new file which differs on only one value?
Absolutely not! Keep in mind that the goal is to create an environment where everything is easily reproducible
(one could suggest overriding hydra config at terminal level just for that one value – we strongly disrecommend it for the same reasons as <code>argpase</code>-ing).
<h1 id="putting-it-all-together">Putting it all together</h1>
Of course, we don’t need to put the whole experimental setup in one big <code>yaml</code> file.
Actually, we do strive for modular configs, in the same fashion as we write our code modular.
Consequently, we would have separate files to configure model, data module, optimizer, and even a tensorboard logger.
Here is an outline of a template we use at LeapMind,it is inspired on
<a href="https://github.com/ashleve/lightning-hydra-template">ashleve</a>’s
theme:
<pre><code>.
├── configs
│ ├── callbacks
│ ├── datamodule
│ ├── experiment
│ ├── logger
│ ├── model
│ ├── optimizer
│ ├── scheduler
│ ├── trainer
│ ├── transforms
│ └── config.yaml
├── src
│ ├── architecture 
│ ├── datasets 
│ ├── pl_datamodules 
│ ├── pl_models
│ └── utils
└── run.py
</code></pre>
Generally, <code>src</code> includes all python scripts and <code>configs</code> has all <code>yaml</code> files.
The heart of the system is <code>run.py</code> – it’s hydra wrapper which glues both <code>configs</code> with <code>src</code>.
See
<a href="https://hydra.cc/docs/intro/#quick-start-guide">official tutorial</a>
to learn how to do it (spoiler: it’s super easy, it’s just one decorator).
We recommend that initially <code>configs</code> contain the default configuration of a model, e.g a baseline.
Once confings for each single part of a project are defined, “aggregate” them in one master file, <code>config.yaml</code>:
<pre><code class="hljs language-yaml">defaults:
 - trainer: superresnet50_trainer.yaml
 - model: superresnet50.yaml
 - optimizer: sgd.yaml
 - datamodule: imagenet.yaml
 - logger: tensorboard.yaml 
</code></pre>
Suppose you need to conduct an experiment with a new lr value.
First, create a new config under <code>configs/experiments</code>, then include the default configs, and finally override only those that you need:
<pre><code class="hljs language-yaml">defaults:
 - override /trainer: superresnet50_trainer.yaml
 - override /model: superresnet50.yaml
 - override /optimizer: sgd.yaml
 - override /datamodule: imagenet.yaml
 - override /logger: tensorboard.yaml

datamodule:
data_dir: /data/imagenet-small
	batch_size: 64

model:
num_classes: 200
learning_rate: 1e-3
</code></pre>
In the development stage, feel free to modify the experiment configs.
However, as soon as you made progress, freeze the configs and share it with your team, e.g by merging it to master branch on git, where it would remain untouched forever.
If you want to add further changes, create a new config file!
<h2 id="ready-to-boost-your-ai-research">Ready to boost your AI research?</h2>
Adding more layers of abstraction to the project might seem just cumbersome – but believe us – it’s really not!
At LeapMind, hydra+lightning template was initially deployed on a small internal project,
but eventually ended up being used company-wide, with many engineers using the same config files.
It did save us many hours and it significantly boosted our workflow.
Are you ready to boost your AI research now?
Or maybe would you like to write some configs <a href="https://leapmind.io/careers/">with us</a>?
 
<hr>
 
Written by Przemyslaw Joniak
Przemyslaw is a graduate student at The University of Tokyo. He appears at LeapMind a few days a week where he works on improving performance of our models.

Learn the best practices we use at LeapMind.

Tech Blog