--- date: '2024-11-11' description: tidbits id: PyTorch modified: 2026-10-01 09:15:44 GMT-04:00 tags: - ml - framework title: PyTorch created: '2024-11-11' published: '2024-11-11' pageLayout: default slug: thoughts/PyTorch permalink: https://aarnphm.xyz/thoughts/PyTorch.md generator: quartz: v4.6.0 hostedProvider: Cloudflare baseUrl: aarnphm.xyz full: https://aarnphm.xyz/llms-full.txt --- see also: [unstable docs](https://pytorch.org/docs/main/) A fixed logit can receive different probabilities in sequences of different lengths. [Softmax](https://docs.pytorch.org/docs/stable/generated/torch.nn.Softmax.html) normalizes over every score: $$ p_i = \frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}}. $$ This puts the same maximum at the start of two random score vectors. Their denominators differ. ```python title="qk_score.py" import torch torch.manual_seed(0) qk_scores_short = torch.randn(2048) qk_scores_long = torch.randn(128000) max_v = torch.max(qk_scores_short.max(), qk_scores_long.max()) qk_scores_short[0] = max_v qk_scores_long[0] = max_v print(qk_scores_short.softmax(0)[0], qk_scores_long.softmax(0)[0]) ``` ## `MultiMarginLoss` [`MultiMarginLoss`](https://docs.pytorch.org/docs/stable/generated/torch.nn.MultiMarginLoss.html) compares the target class score with each competing score. For one example, let $x \in \mathbb{R}^C$ and let $y \in \{0,\ldots,C-1\}$ be the target class index. With no class weights, $$ \ell(x,y) = \frac{1}{C}\sum_{i \ne y}\max(0, m - x_y + x_i)^p, \qquad p \in \{1,2\}. $$ A competing class contributes zero once the target score exceeds it by at least the margin $m$. The default uses $m=1$, $p=1$, and averages the per-example losses over the batch. The division by $C$ still includes the target class in the class count, even though its term is excluded from the sum. ## `SGD` PyTorch’s [SGD implementation](https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html) applies the `maximize` sign to the objective gradient before adding weight decay. The parameter update then subtracts the resulting direction in both modes. This keeps decay pointing toward zero. For [[thoughts/Nesterov momentum]], PyTorch cites [On the importance of initialization and momentum in deep learning](http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf). Its momentum buffer is initialized to the first gradient. Dampening starts at the second update; Nesterov requires positive momentum and zero dampening.
"\\begin{algorithm}\n\\caption{SGD in PyTorch}\n\\begin{algorithmic}\n\\State \\textbf{input:} $\\gamma$ (lr), $\\theta_0$ (params), $f(\\theta)$ (objective), $\\lambda$ (weight decay),\n\\State $\\mu$ (momentum), $\\tau$ (dampening), nesterov, maximize\n\\For{$t = 1$ to $...$}\n \\State $g_t \\gets \\nabla_\\theta f_t(\\theta_{t-1})$\n \\If{$\\text{maximize}$}\n \\State $g_t \\gets -g_t$\n \\EndIf\n \\If{$\\lambda \\neq 0$}\n \\State $g_t \\gets g_t + \\lambda\\theta_{t-1}$\n \\EndIf\n \\If{$\\mu \\neq 0$}\n \\If{$t > 1$}\n \\State $b_t \\gets \\mu b_{t-1} + (1-\\tau)g_t$\n \\Else\n \\State $b_t \\gets g_t$\n \\EndIf\n \\If{$\\text{nesterov}$}\n \\State $g_t \\gets g_t + \\mu b_t$\n \\Else\n \\State $g_t \\gets b_t$\n \\EndIf\n \\EndIf\n \\State $\\theta_t \\gets \\theta_{t-1} - \\gamma g_t$\n\\EndFor\n\\State \\textbf{return} $\\theta_t$\n\\end{algorithmic}\n\\end{algorithm}"

Algorithm 1 SGD in PyTorch

1:input: γ\gamma (lr), θ0\theta_0 (params), f(θ)f(\theta) (objective), λ\lambda (weight decay),

2:μ\mu (momentum), τ\tau (dampening), nesterov, maximize

3:for t=1t = 1 to ...... do

4:gt←∇θft(θt−1)g_t \gets \nabla_\theta f_t(\theta_{t-1})

5:if maximize\text{maximize} then

6:gt←−gtg_t \gets -g_t

7:end if

8:if λ≠0\lambda \neq 0 then

9:gt←gt+λθt−1g_t \gets g_t + \lambda\theta_{t-1}

10:end if

11:if μ≠0\mu \neq 0 then

12:if t>1t > 1 then

13:bt←μbt−1+(1−τ)gtb_t \gets \mu b_{t-1} + (1-\tau)g_t

14:else

15:bt←gtb_t \gets g_t

16:end if

17:if nesterov\text{nesterov} then

18:gt←gt+μbtg_t \gets g_t + \mu b_t

19:else

20:gt←btg_t \gets b_t

21:end if

22:end if

23:θt←θt−1−γgt\theta_t \gets \theta_{t-1} - \gamma g_t

24:end for

25:return θt\theta_t

## [[thoughts/knowledge distillation]] CIFAR-10 examples from the [PyTorch distillation tutorial](https://docs.pytorch.org/tutorials/beginner/knowledge_distillation_tutorial.html). These are excerpts: the loaders, trained teacher, and model instances must be set up before running the training loops. ```python title="distill.py" path="thoughts/scripts/distill.py" import torch import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms import torchvision.datasets as datasets # Check if the current `accelerator `__ # is available, and if not, use the CPU device = ( torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else 'cpu' ) print(f'Using {device} device') # Below we are preprocessing data for CIFAR-10. We use an arbitrary batch size of 128. transforms_cifar = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # Loading the CIFAR-10 dataset: train_dataset = datasets.CIFAR10( root='./data', train=True, download=True, transform=transforms_cifar ) test_dataset = datasets.CIFAR10( root='./data', train=False, download=True, transform=transforms_cifar ) # Deeper neural network class to be used as teacher: class DeepNN(nn.Module): def __init__(self, num_classes=10): super(DeepNN, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 128, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(128, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(64, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Linear(2048, 512), nn.ReLU(), nn.Dropout(0.1), nn.Linear(512, num_classes), ) def forward(self, x): x = self.features(x) x = torch.flatten(x, 1) x = self.classifier(x) return x # Lightweight neural network class to be used as student: class LightNN(nn.Module): def __init__(self, num_classes=10): super(LightNN, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(16, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Linear(1024, 256), nn.ReLU(), nn.Dropout(0.1), nn.Linear(256, num_classes), ) def forward(self, x): x = self.features(x) x = torch.flatten(x, 1) x = self.classifier(x) return x def train(model, train_loader, epochs, learning_rate, device): criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) model.train() for epoch in range(epochs): running_loss = 0.0 for inputs, labels in train_loader: # inputs: A collection of batch_size images # labels: A vector of dimensionality batch_size with integers denoting class of each image inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) # outputs: Output of the network for the collection of images. A tensor of dimensionality batch_size x num_classes # labels: The actual labels of the images. Vector of dimensionality batch_size loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print( f'Epoch {epoch + 1}/{epochs}, Loss: {running_loss / len(train_loader)}' ) def test(model, test_loader, device): model.to(device) model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total print(f'Test Accuracy: {accuracy:.2f}%') return accuracy def train_knowledge_distillation( teacher, student, train_loader, epochs, learning_rate, T, soft_target_loss_weight, ce_loss_weight, device, ): ce_loss = nn.CrossEntropyLoss() optimizer = optim.Adam(student.parameters(), lr=learning_rate) teacher.eval() # Teacher set to evaluation mode student.train() # Student to train mode for epoch in range(epochs): running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() # Forward pass with the teacher model - do not save gradients here as we do not change the teacher's weights with torch.no_grad(): teacher_logits = teacher(inputs) # Forward pass with the student model student_logits = student(inputs) # Soften the student logits by applying softmax first and log() second soft_targets = nn.functional.softmax(teacher_logits / T, dim=-1) soft_prob = nn.functional.log_softmax(student_logits / T, dim=-1) # Calculate the soft targets loss. Scaled by T**2 as suggested by the authors of the paper "Distilling the knowledge in a neural network" soft_targets_loss = ( torch.sum(soft_targets * (soft_targets.log() - soft_prob)) / soft_prob.size()[0] * (T**2) ) # Calculate the true label loss label_loss = ce_loss(student_logits, labels) # Weighted sum of the two losses loss = ( soft_target_loss_weight * soft_targets_loss + ce_loss_weight * label_loss ) loss.backward() optimizer.step() running_loss += loss.item() print( f'Epoch {epoch + 1}/{epochs}, Loss: {running_loss / len(train_loader)}' ) if __name__ == '__main__': torch.manual_seed(42) nn_deep = DeepNN(num_classes=10).to(device) train(nn_deep, train_loader, epochs=10, learning_rate=0.001, device=device) test_accuracy_deep = test(nn_deep, test_loader, device) # Instantiate the lightweight network: torch.manual_seed(42) nn_light = LightNN(num_classes=10).to(device) print( 'Norm of 1st layer of nn_light:', torch.norm(nn_light.features[0].weight).item(), ) print( 'Norm of 1st layer of new_nn_light:', torch.norm(new_nn_light.features[0].weight).item(), ) train(nn_light, train_loader, epochs=10, learning_rate=0.001, device=device) test_accuracy_light_ce = test(nn_light, test_loader, device) print(f'Teacher accuracy: {test_accuracy_deep:.2f}%') print(f'Student accuracy: {test_accuracy_light_ce:.2f}%') # Apply ``train_knowledge_distillation`` with a temperature of 2. Arbitrarily set the weights to 0.75 for CE and 0.25 for distillation loss. train_knowledge_distillation( teacher=nn_deep, student=new_nn_light, train_loader=train_loader, epochs=10, learning_rate=0.001, T=2, soft_target_loss_weight=0.25, ce_loss_weight=0.75, device=device, ) test_accuracy_light_ce_and_kd = test(new_nn_light, test_loader, device) # Compare the student test accuracy with and without the teacher, after distillation print(f'Teacher accuracy: {test_accuracy_deep:.2f}%') print(f'Student accuracy without teacher: {test_accuracy_light_ce:.2f}%') print(f'Student accuracy with CE + KD: {test_accuracy_light_ce_and_kd:.2f}%') ``` ### Cosine loss minimisation run The experiment asks whether a trained teacher’s hidden [[thoughts/representations]] help the student classify images. The teacher stays fixed. The student minimizes a weighted sum of classification loss and cosine loss between their hidden activations. [`CosineEmbeddingLoss`](https://docs.pytorch.org/docs/stable/generated/torch.nn.CosineEmbeddingLoss.html) compares equal-length vectors in their current coordinate order. It encourages aligned directions. It does not copy weights or find a permutation of the teacher’s features: $$ \text{loss}(x,y) = \begin{cases} 1 - \cos(x_1, x_2), & \text{if } y = 1 \\ \max(0, \cos(x_1, x_2) - \text{margin}), & \text{if } y = -1 \end{cases} $$ Here the loss target is $y=1$ for every teacher/student pair, independently of the image’s class label. Positive rescaling leaves the cosine unchanged, so this objective does not require equal activation magnitudes. The teacher produces $2048$ flattened features. Average pooling reduces these to $1024$, matching the student’s $1024$ features.[^internal] That makes the loss computable; whether this alignment helps classification still needs a held-out comparison. The cosine term updates the student’s feature extractor, while cross-entropy also trains its classifier. [^internal]: Both modified models return logits and the hidden representation: ```python sample_input = torch.randn(128, 3, 32, 32).to( device ) # Batch size: 128, Channels: 3, Image size: 32x32 logits, hidden_representation = modified_nn_light(sample_input) print('Student logits shape:', logits.shape) # batch_size x total_classes print( 'Student hidden representation shape:', hidden_representation.shape ) # batch_size x hidden_representation_size logits, hidden_representation = modified_nn_deep(sample_input) print('Teacher logits shape:', logits.shape) # batch_size x total_classes print( 'Teacher hidden representation shape:', hidden_representation.shape ) # batch_size x hidden_representation_size ``` ```python title="modified_deep_cosine.py" path="thoughts/scripts/modified_deep_cosine.py" class ModifiedDeepNNCosine(nn.Module): def __init__(self, num_classes=10): super(ModifiedDeepNNCosine, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 128, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(128, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(64, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Linear(2048, 512), nn.ReLU(), nn.Dropout(0.1), nn.Linear(512, num_classes), ) def forward(self, x): x = self.features(x) flattened_conv_output = torch.flatten(x, 1) x = self.classifier(flattened_conv_output) flattened_conv_output_after_pooling = torch.nn.functional.avg_pool1d( flattened_conv_output, 2 ) return x, flattened_conv_output_after_pooling # Create a similar student class where we return a tuple. We do not apply pooling after flattening. class ModifiedLightNNCosine(nn.Module): def __init__(self, num_classes=10): super(ModifiedLightNNCosine, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(16, 16, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Linear(1024, 256), nn.ReLU(), nn.Dropout(0.1), nn.Linear(256, num_classes), ) def forward(self, x): x = self.features(x) flattened_conv_output = torch.flatten(x, 1) x = self.classifier(flattened_conv_output) return x, flattened_conv_output def train_cosine_loss( teacher, student, train_loader, epochs, learning_rate, hidden_rep_loss_weight, ce_loss_weight, device, ): ce_loss = nn.CrossEntropyLoss() cosine_loss = nn.CosineEmbeddingLoss() optimizer = optim.Adam(student.parameters(), lr=learning_rate) teacher.to(device) student.to(device) teacher.eval() # Teacher set to evaluation mode student.train() # Student to train mode for epoch in range(epochs): running_loss = 0.0 for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() # Forward pass with the teacher model and keep only the hidden representation with torch.no_grad(): _, teacher_hidden_representation = teacher(inputs) # Forward pass with the student model student_logits, student_hidden_representation = student(inputs) # Calculate the cosine loss. Target is a vector of ones. From the loss formula above we can see that is the case where loss minimization leads to cosine similarity increase. hidden_rep_loss = cosine_loss( student_hidden_representation, teacher_hidden_representation, target=torch.ones(inputs.size(0)).to(device), ) # Calculate the true label loss label_loss = ce_loss(student_logits, labels) # Weighted sum of the two losses loss = ( hidden_rep_loss_weight * hidden_rep_loss + ce_loss_weight * label_loss ) loss.backward() optimizer.step() running_loss += loss.item() print( f'Epoch {epoch + 1}/{epochs}, Loss: {running_loss / len(train_loader)}' ) def test_multiple_outputs(model, test_loader, device): model.to(device) model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs, _ = model(inputs) # Disregard the second tensor of the tuple _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total print(f'Test Accuracy: {accuracy:.2f}%') return accuracy if __name__ == '__main__': # We do not have to train the modified deep network from scratch of course, we just load its weights from the trained instance modified_nn_deep = ModifiedDeepNNCosine(num_classes=10).to(device) modified_nn_deep.load_state_dict(nn_deep.state_dict()) # Once again ensure the norm of the first layer is the same for both networks print( 'Norm of 1st layer for deep_nn:', torch.norm(nn_deep.features[0].weight).item(), ) print( 'Norm of 1st layer for modified_deep_nn:', torch.norm(modified_nn_deep.features[0].weight).item(), ) # Initialize a modified lightweight network with the same seed as our other lightweight instances. This will be trained from scratch to examine the effectiveness of cosine loss minimization. torch.manual_seed(42) modified_nn_light = ModifiedLightNNCosine(num_classes=10).to(device) print( 'Norm of 1st layer:', torch.norm(modified_nn_light.features[0].weight).item(), ) ```