Use case for multiple calls to `apex.amp.initialize`

Hello,

The documentation for the apex.amp.initialize function states, "If you think your use case requires amp.initialize to be called more than once, let us know."

My use case involves a neural architecture search technique called PC-DARTS, detailed in this paper and implemented in this repository. The implementation uses two optimizers and results in the dynamic creation of many models (one new model is created for each batch of training data that is loaded). The same use case would apply for the original DARTS implementation that PC-DARTS is based on.

In particular, running python train_search.py --unrolled on the train_search.py file found in the PC-DARTS repository results in the following situation:

• A "main" model is instantiated before training begins, in train_search.py#L77.
• Two optimizers are instantiated before training begins, in train_search.py#L81 (torch.optim.SGD) and in architect.py#L17 (torch.optim.Adam). The constructor for each of these optimizers is called on a different subset of the parameters from the main model.
• For each batch of data loaded during the training loop, the following steps occur:
  • A new model is created at architect.py#L62; it is instantiated with weights similar to (but not identical to) those used in the main model.
  • a loss is calculated via a forward pass through this newly-created model; backpropagation occurs at architect.py#L47. Some of the backpropagated gradients from the newly-created model are then copied over to the main model (at architect.py#L55). The newly-created model is discarded shortly afterwards.
  • next, an optimizer step is taken at architect.py#L37 using one of the two optimizers. This results in an update to the weights of the main model. The step is taken based on the gradients that were copied over to the main model from the dynamically-created model.
  • finally, another loss is calculated, this time via a forward pass through the "main" model. Backpropagation occurs at train_search.py#L163 and an optimizer step is taken at train_search.py#L165 using the other of the two optimizers. This again results in an update to the main model's weights.

If possible, I would like to use mixed precision for both the main model as well as for computation conducted in the dynamically-created models.

Start with the apex.amp.initialize documentation and the referenced PC-DARTS entry points: train_search.py and architect.py. Trace the two optimizers, the dynamically created models, and both optimizer steps to determine the requirements for mixed precision in this workflow. Done means the requested use case is supported or its limitations and required initialization behavior are documented.