ptorch常用代码梯度篇（梯度裁剪、梯度累积、冻结预训练层等）

梯度裁剪（Gradient Clipping）

在训练比较深或者循环神经网络模型的过程中，我们有可能发生梯度爆炸的情况，这样会导致我们模型训练无法收敛。 我们可以采取一个简单的策略来避免梯度的爆炸，那就是梯度截断 Clip, 将梯度约束在某一个区间之内，在训练的过程中，在优化器更新之前进行梯度截断操作。!!!!! 注意这个方法只在训练的时候使用，在测试的时候验证和测试的时候不用。

整个流程简单总结如下：

1. 加载训练数据和标签
2. 模型输入输出
3. 计算 loss 函数值
4. loss 反向传播
5. 梯度截断
6. 优化器更新梯度参数

import torch.nn as nn
outputs = model(data)
loss= loss_fn(outputs, target)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=20, norm_type=2)
optimizer.step()
optimizer.zero_grad()

nn.utils.clip_grad_norm_ 输入是（NN 参数，最大梯度范数，范数类型 = 2) 一般默认为 L2 范数。

梯度累积

常规网络如下：

# 正常网络
optimizer.zero_grad()
for idx, (x, y) in enumerate(train_loader):
    pred = model(x)
    loss = criterion(pred, y)
    
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    
    if (idx+1) % eval_steps == 0:
        eval()

需要梯度累计时，每个 mini-batch 仍然正常前向传播以及反向传播，但是反向传播之后并不进行梯度清零，因为 PyTorch 中的 loss.backward() 执行的是梯度累加的操作，所以当我们调用 4 次 loss.backward() 后，这 4 个 mini-batch 的梯度都会累加起来。但是，我们需要的是一个平均的梯度，或者说平均的损失，所以我们应该将每次计算得到的 loss除以 accum_steps

# 梯度累积

accum_steps = 4
optimizer.zero_grad()
for idx, (x, y) in enumerate(train_loader):
    pred = model(x)
    loss = criterion(pred, y)
    
    # normlize loss to account for batch accumulation
    loss = loss / accum_steps
    
    loss.backward()
    
    if (idx+1) % accum_steps == 0 or (idx+1) == len(train_loader):
        optimizer.step()
        optimizer.zero_grad()
    if (idx+1) % eval_steps == 0:
            eval()

总的来说，梯度累加就是计算完每个 mini-batch 的梯度后不清零，而是做梯度的累加，当累加到一定的次数之后再更新网络参数，然后将梯度清零。通过这种延迟更新的手段，可以实现与采用大 batch_size 相近的效果

冻结某些层

在加载预训练模型的时候，我们有时想冻结前面几层，使其参数在训练过程中不发生变化。

def freeze(module):
    """
    Freezes module's parameters.
    """
    
    for parameter in module.parameters():
        parameter.requires_grad = False
        
def get_freezed_parameters(module):
    """
    Returns names of freezed parameters of the given module.
    """
    
    freezed_parameters = []
    for name, parameter in module.named_parameters():
        if not parameter.requires_grad:
            freezed_parameters.append(name)
            
    return freezed_parameters

import torch
from transformers import AutoConfig, AutoModel


# initializing model
model_path = "microsoft/deberta-v3-base"
config = AutoConfig.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, config=config)


# freezing embeddings and first 2 layers of encoder
freeze(model.embeddings)
freeze(model.encoder.layer[:2])

freezed_parameters = get_freezed_parameters(model)
print(f"Freezed parameters: {freezed_parameters}")

# selecting parameters, which requires gradients and initializing optimizer
model_parameters = filter(lambda parameter: parameter.requires_grad, model.parameters())
optimizer = torch.optim.AdamW(params=model_parameters, lr=2e-5, weight_decay=0.0)

Freezed parameters: ['embeddings.word_embeddings.weight', 'embeddings.LayerNorm.weight', 'embeddings.LayerNorm.bias', 'encoder.layer.0.attention.self.query_proj.weight', 'encoder.layer.0.attention.self.query_proj.bias', 'encoder.layer.0.attention.self.key_proj.weight', 'encoder.layer.0.attention.self.key_proj.bias', 'encoder.layer.0.attention.self.value_proj.weight', 'encoder.layer.0.attention.self.value_proj.bias', 'encoder.layer.0.attention.output.dense.weight', 'encoder.layer.0.attention.output.dense.bias', 'encoder.layer.0.attention.output.LayerNorm.weight', 'encoder.layer.0.attention.output.LayerNorm.bias', 'encoder.layer.0.intermediate.dense.weight', 'encoder.layer.0.intermediate.dense.bias', 'encoder.layer.0.output.dense.weight', 'encoder.layer.0.output.dense.bias', 'encoder.layer.0.output.LayerNorm.weight', 'encoder.layer.0.output.LayerNorm.bias', 'encoder.layer.1.attention.self.query_proj.weight', 'encoder.layer.1.attention.self.query_proj.bias', 'encoder.layer.1.attention.self.key_proj.weight', 'encoder.layer.1.attention.self.key_proj.bias', 'encoder.layer.1.attention.self.value_proj.weight', 'encoder.layer.1.attention.self.value_proj.bias', 'encoder.layer.1.attention.output.dense.weight', 'encoder.layer.1.attention.output.dense.bias', 'encoder.layer.1.attention.output.LayerNorm.weight', 'encoder.layer.1.attention.output.LayerNorm.bias', 'encoder.layer.1.intermediate.dense.weight', 'encoder.layer.1.intermediate.dense.bias', 'encoder.layer.1.output.dense.weight', 'encoder.layer.1.output.dense.bias', 'encoder.layer.1.output.LayerNorm.weight', 'encoder.layer.1.output.LayerNorm.bias']

可以看到前两层的 weight 和 bias 的 requires_grad 都为 False，表示它们不可训练。

最后在定义优化器时，只对 requires_grad 为 True 的层的参数进行更新。（这里用filter筛选只传入了requires_grad为True的参数，但如果直接传入全部参数应该也可以达到只训练未冻结层参数的效果）

optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=0.01)

其他注意事项

1. with torch.no_grad()或者@torch.no_grad()中的数据不需要计算梯度，也不会进行反向传播。不需要计算梯度的代码块(如验证测试)用 with torch.no_grad() 包含起来，节省显存

model.eval()
with torch.no_grad():
   pass

@torch.no_grad()
def eval():
	...

2. model.eval() 和 torch.no_grad() 的区别在于，model.eval() 是将网络切换为测试状态，例如 BN 和dropout在训练和测试阶段使用不同的计算方法。torch.no_grad() 是关闭 PyTorch 张量的自动求导机制，以减少存储使用和加速计算，得到的结果无法进行 loss.backward()。

3. model.zero_grad()会把整个模型的参数的梯度都归零, 而optimizer.zero_grad()只会把传入其中的参数的梯度归零.

4. loss.backward() 前用 optimizer.zero_grad() 清除累积梯度。如果在循环里需要把optimizer.zero_grad()写在后面，那应该在循环外需要先调用一次optimizer.zero_grad()

5. 查看网络中的梯度

params = list(model.named_parameters())
(name, param) = params[28]
print(name)
print(param.grad)
print('-------------------------------------------------')