Transfer learning with mixed-precision

username: Sayak_Paul
cooked: <p>Hi.</p>
<p>I am currently trying to figure out how to facilitate mixed-precision training when using transfer learning.</p>
<p>Typically, here’s how I am building my custom model using a pre-trained base model:</p>
<pre><code class="lang-python"># load up the ResNet50 model
model = models.resnet50(pretrained=True)

# append a new classification top to our feature extractor and pop it
# on to the current device
modelOutputFeats = model.fc.in_features
model.fc = nn.Linear(modelOutputFeats, len(trainDS.classes))
model = model.to(DEVICE)
</code></pre>
<p>I am also using <code>nn.DataParallel</code>. As demonstrated <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#dataparallel-in-a-single-process" rel="noopener nofollow ugc">here</a>, simply putting <code>model</code> inside <code>autocast</code> context, won’t help. So, is there a way I could perform it using the practices recommended <a href="https://pytorch.org/docs/stable/notes/amp_examples.html" rel="noopener nofollow ugc">here</a>?</p>

username: Sayak_Paul
cooked: <p>Doing the following is working for me but I am still a bit unsure if the <code>autocast</code> state if being sent to multiple devices. <a class="mention" href="/u/ptrblck">@ptrblck</a> could you advise?</p>
<pre><code class="lang-python">for (i, (x, y)) in enumerate(trainLoader):
	with torch.cuda.amp.autocast(enabled=True):
		# send the input to the device
		(x, y) = (x.to(DEVICE), y.to(DEVICE))

		# perform a forward pass and calculate the training loss
		pred = model(x)
		loss = lossFunc(pred, y)

	# calculate the gradients
	scaler.scale(loss).backward()
	scaler.step(opt)
	scaler.update()
	opt.zero_grad()
</code></pre>

username: ptrblck
cooked: <p>I’m not sure I understand the issue correctly. Could you explain what <code>autocast</code> state would refer to and what the initial issue was that you were seeing?<br>
The linked docs give an example of how to use <code>nn.DataParallel</code>. Are you seeing any errors with it?</p>

username: Sayak_Paul
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> the linked doc (first one) discusses AMP when using <code>nn.DataParallel</code>. It states that the <code>forward</code> method of my model needs to be inside an <code>autocast</code> context.</p>
<p>My question is how do I achieve that when I am using a pre-trained model? Would the snippet I provided work in that case?</p>
<p>The <code>autocast</code> state would refer to the <code>with torch.cuda.amp.autocast(enabled=True):</code> context.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Sayak_Paul" data-post="4" data-topic="117350">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sayak_paul/48/12743_2.png" class="avatar"> Sayak_Paul:</div>
<blockquote>
<p>My question is how do I achieve that when I am using a pre-trained model?</p>
</blockquote>
</aside>
<p>You could write a custom model, decorate the <code>forward</code> method of the custom model with the <code>autocast</code> context, and call the pretrained model inside.</p>

username: Sayak_Paul
cooked: <p>This is what I am thinking as well but I am not clear on how would I only extract the backbone of a pre-trained model. Mentally, I have the following <code>forward</code> pass:</p>
<pre><code class="lang-auto">def forward (self, x)
    with autocast(enabled=True):
        features = self.base_model(x)
        logits = self.classifier(features)
    return logits
</code></pre>
<p>My question is how should I extract the <code>base_model</code> from let’s say, <code>torchvision.models.resnet50(pretrained=True</code> throwing out the <code>fc</code> block from it.</p>

username: ptrblck
cooked: <p>Using the backbone of a pretrained model seems to be orthogonal to the usage of amp and you could call the model directly in your custom model wrapper.</p>
<p>However, if you want to only use the backbone, replace other layers either with <code>nn.Identity</code> layers or check the source code of the model and see, if you could simply call the backbone directly.<br>
Some models are implemented in this way, others are more complicated and would need to a bit more attention.</p>

username: Sayak_Paul
cooked: <p>Okay. Thank you for the points. Appreciate the help.</p>

username: ptrblck
cooked: <p>Since you are working with a <code>ResNet</code>, have a look at <a href="https://github.com/pytorch/vision/blob/e79a74e1bfdf4e8e0275eeecea9063adba600ba0/torchvision/models/resnet.py#L230-L246">its <code>forward</code></a>. As you can see, there is unfortunately no <code>model.features</code> attribute, so in this case it would be easier to replace <code>model.classifier = nn.Identity()</code> and just call the model directly.</p>

username: Sayak_Paul
cooked: <p>So, here’s what I did:</p>
<pre><code class="lang-python">from torch.cuda.amp import autocast
from torch import nn

class CustomClassifier(nn.Module):
	def __init__(self, baseModel, numClasses):
		super(FoodClassier, self).__init__()

		# initialize the base model and the classification layer
		self.baseModel = baseModel
		self.classifier = nn.Linear(baseModel.classifier.in_features,
			numClasses)

		# set the classifier of our base model to produce outputs
		# from the last convolution block
		self.baseModel.classifier = nn.Identity()

	# we decorate the `forward()` method with `autocast()` to enable
	# mixed-precision training in a distributed manner
	@autocast(enabled=True)
	def forward(self, x):
		# pass the inputs through the base model and then obtain the
		# obtain the classifier outputs
		features = self.baseModel(x)
		logits = self.classifier(features)

		# return the classifier outputs
		return logits
</code></pre>
<p>I am using a <code>densenet121</code> hence changed it to <code>classifier</code> from <code>fc</code>. However, I don’t see any speed-up in the total model training after making these changes.</p>