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,
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.
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
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.
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")