NAS Quick Start¶
The NAS feature provided by NNI has two key components: APIs for expressing the search space and NAS training approaches. The former is for users to easily specify a class of models (i.e., the candidate models specified by the search space) which may perform well. The latter is for users to easily apply state-of-the-art NAS training approaches on their own model.
Here we use a simple example to demonstrate how to tune your model architecture with the NNI NAS APIs step by step. The complete code of this example can be found here.
Write your model with NAS APIs¶
Instead of writing a concrete neural model, you can write a class of neural models using two of the NAS APIs library functions, LayerChoice
and InputChoice
. For example, if you think either of two options might work in the first convolution layer, then you can get one from them using LayerChoice
as shown by self.conv1
in the code. Similarly, the second convolution layer self.conv2
also chooses one from two options. To this line, four candidate neural networks are specified. self.skipconnect
uses InputChoice
to specify two choices, adding a skip connection or not.
import torch.nn as nn
from nni.nas.pytorch.mutables import LayerChoice, InputChoice
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = LayerChoice([nn.Conv2d(3, 6, 3, padding=1), nn.Conv2d(3, 6, 5, padding=2)])
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = LayerChoice([nn.Conv2d(6, 16, 3, padding=1), nn.Conv2d(6, 16, 5, padding=2)])
self.conv3 = nn.Conv2d(16, 16, 1)
self.skipconnect = InputChoice(n_candidates=1)
self.bn = nn.BatchNorm2d(16)
self.gap = nn.AdaptiveAvgPool2d(4)
self.fc1 = nn.Linear(16 * 4 * 4, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
For a detailed description of LayerChoice
and InputChoice
, please refer to the NAS guide
Choose a NAS trainer¶
After the model is instantiated, it is time to train the model using a NAS trainer. Different trainers use different approaches to search for the best one from a class of neural models that you specified. NNI provides several popular NAS training approaches such as DARTS and ENAS. Here we use DartsTrainer
in the example below. After the trainer is instantiated, invoke trainer.train()
to do the search.
trainer = DartsTrainer(net,
loss=criterion,
metrics=accuracy,
optimizer=optimizer,
num_epochs=2,
dataset_train=dataset_train,
dataset_valid=dataset_valid,
batch_size=64,
log_frequency=10)
trainer.train()
Export the best model¶
After the search (i.e., trainer.train()
) is done, to get the best performing model we simply call trainer.export("final_arch.json")
to export the found neural architecture to a file.
NAS visualization¶
We are working on NAS visualization and will release this feature soon.
Retrain the exported best model¶
It is simple to retrain the found (exported) neural architecture. Step one, instantiate the model you defined above. Step two, invoke apply_fixed_architecture
to the model. Then the model becomes the found (exported) one. Afterward, you can use traditional training to train this model.
model = Net()
apply_fixed_architecture(model, "final_arch.json")