no_grad() is a PyTorch context manager for computations that should not produce gradients. Wrap prediction or evaluation code in with torch.no_grad(): to avoid the extra time and memory of gradient calculation and reduce the risk of interfering with later training. section 2.1.7 section 2.4.4
Use the context when:
- Making predictions or running inference after training. section 2.1.7
- Evaluating a model on held-out data after switching it to model.eval(). section 2.4.4
- Freezing a pretrained component such as BERT while training a new classifier around it. section 13.4.2
Prediction without gradient tracking
with torch.no_grad(): Y_pred = model(torch.tensor(X.reshape(-1, 1), device=device, dtype=torch.float32)).cpu().numpy() The model produces predictions for plotting, while PyTorch avoids recording the computation for backpropagation. section 2.1.7
Evaluation and generation
The same pattern evaluates a trained GAN: sample noise, generate fake digits, score them with the discriminator, and move results to the CPU without building a gradient graph. Seq2Seq decoding also uses the block while selecting next words, then exits it before later operations that need gradients. section 9.1.3 section 11.4.1
class BertBasedClassifier(nn.Module): #1
def __init__(self, bert_model, classes):
"""
bert_model: the BERT-based classification model to use as a frozen
initial layer of the network
classes: the number of output neurons/target classes for this
classifier.
"""
super(BertBasedClassifier, self).__init__()
self.bert_model = bert_model #2
self.attn = AttentionAvg( #3
AdditiveAttentionScore(
bert_model.config.dim))
self.fc1 = nn.Linear(bert_model.config.dim, #4
bert_model.config.dim)
self.pred = nn.Linear(bert_model.config.dim, classes) #5
def forward(self, input):
mask = getMaskByFill(input) #6
with torch.no_grad(): #7
x = self.bert_model(input)[0] #8
cntxt = x.sum(dim=1)/(mask.sum(dim=1).unsqueeze(1)+1e-5) #9
x = self.attn(x, cntxt, mask) #10
x = F.relu(self.fc1(x)) #11
return self.pred(x)
bertClassifier = BertBasedClassifier(bert_model, NUN_CLASS) #12
bert_results = train_network(bertClassifier, nn.CrossEntropyLoss(),
train_text_bert_loader, test_loader=test_text_bert_loader,
device=device, epochs=10, score_funcs={’Accuracy’: accuracy_score})
#1 Our new class for frozen training of BERT models
#2 We get a tensor of shape (B, T, D) from BERT, so we define a few of our own layers to get from (B, T, D) to a prediction of shape (B, classes)
#3 Attention to get down to (B, D) shape
#4 Does a little feature extraction
#5 Makes a prediction about the class
#6 Input is (B, T).
#7 This with no_grad() does the freezing.
#8 Hugging Face returns a tuple, so unpack it! (B, T, D)
#9 Computes the average embedding
#10 Applies the attention
#11 Makes predictions and returns
#12 Builds the classifier
Shows no_grad() inside a BERT-based classifier, where pretrained representations feed trainable classification layers. section 13.4.2
