Save and load a mixed precision model

Rafi: <p>I am training distilBert model for text classification. I am using the amp package to train the mixed precision version. here is the code to train the model</p>
<pre><code class="lang-auto">     # define the model
      model = BertMulticlassifier(....)

    # Define optimizer and scheduler
    optimizer = Config.OPTIMIZER(model.parameters(), lr=Config.LR, bias_correction=False)
    # model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
    scheduler = Config.SCHEDULER(optimizer=optimizer,
                                 num_warmup_steps=400,
                                 num_training_steps=len(train_data_loader) * Config.EPOCHS
                                 )

    model.compile(optimizer=optimizer, loss=loss_func, metrics=metric, scheduler=scheduler)
    model.fit(train_loader=train_data_loader, val_loader=val_data_loader, epochs=Config.EPOCHS)
   # save the model
    torch.save(self.state_dict(), model_name)
</code></pre>
<p>Load the model and run it on CPU:</p>
<pre><code class="lang-auto">model = BertMulticlassifier(....)
model.load_state_dict(torch.load('model.bin', map_location='cpu'))
model.eval()
model.predict()
</code></pre>
<p>By comparing the inference time between the distilbert and the mixed precision distilbert version, it was more or less the same, which does not make sense. I expect that the mixed precision version take less time.</p>
<p>So, I am wondering that the amp just used the mixed precision to make the train and then it saved the model with float32 or the problem is coming from the load function ? What is the right way to  load the model in the mixed precision version. Thanks to take a look <a class="mention" href="/u/ptrblck">@ptrblck</a></p>

ptrblck: <p>It seems you are using a higher-level API, which provides a <code>fit</code> and <code>predict</code> function, and which might use automatic mixed precision internally.<br>
Which library are you using and how did you specify to use amp?</p>

Rafi: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>. Basically, I implemented my own class following the keras schema, using the transformers hugging face package<br>
here is the internal part of the code in fit() function for one epoch</p>
<pre><code class="lang-auto">losses = []
predictions = []
 for step, batch in  enumerate(data_loader):
            input_ids = batch["input_ids"].to(self.device)
            attention_mask = batch["attention_mask"].to(self.device)
            targets = batch["targets"].to(self.device)
            outputs = self.__call__(
                input_ids=input_ids,
                attention_mask=attention_mask
            )
            loss = self.loss_func(outputs, targets)
           ......
           with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                     scaled_loss.backward()
            nn.utils.clip_grad_norm_(self.parameters(), max_norm=1.0)
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()
return np.mean(predictions), np.mean(losses)
</code></pre>
<p>Alos pytorch provide an other way to quantize the model.  I have used this with bertdistill trained with mixed precision. The quantized model performs as good as the float32 version!</p>
<p><em>So my though is that the mixed precision is used just for the training phase and the model is saved all as 32 float??</em></p>
<p><code>model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)</code></p>

ptrblck: <aside class="quote no-group" data-username="Rafi" data-post="3" data-topic="94724">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/779978/48.png" class="avatar"> Rafi:</div>
<blockquote>
<p>So my though is that the mixed precision is used just for the training phase and the model is saved all as 32 float??</p>
</blockquote>
</aside>
<p>Yes, the parameters and buffers would be stored in <code>float32</code>. If you don’t use <code>autocast</code> during inference, then you would execute the standard FP32 model.</p>
<p>Note that mixed precision utilities are now usable in PyTorch directly without <code>apex</code>, so I would recommend to update the code to <a href="https://pytorch.org/docs/stable/amp.html"><code>torch.cuda.amp</code></a>.</p>

Rafi: <p>Thanks so much <a class="mention" href="/u/ptrblck">@ptrblck</a>. I have changed my code using autocast for torch.cuda.amp. However I remark that before with</p>
<pre><code class="lang-auto">model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
</code></pre>
<p>I can use a batch_size of 16, however using the autocast, I could not use more than 8 !!! It seems some operations are not casted to float16 as in apex! I think the autocast implement the <code>level 01</code> of apex</p>

ptrblck: <p>The native amp implementation is corresponding to <code>opt_level="O1"</code> in <code>apex.amp</code>, as the usability is the most user-friendly and it should support all use cases. We are experimenting with other utilities to lower the memory usage, but for now we strongly recommend to stick to the native implementation for the best support.</p>