Способ 1: с помощью оптимизатора PyTorch

Как и широко используемые оптимизаторы, классы Adam и AdamW, предоставляемые инфраструктурой PyTorch, имеют собственный параметр под названием «распад веса», который сам по себе обеспечивает быстрый и простой способ применить регуляризацию L2, просто добавив новый параметр.

Это происходит так: при создании оптимизатора добавьте параметр `weight_decay`.

import torch
optimizer=torch.optim.Adam(model.parameters(),weight_decay=0.01)

При добавлении weight_decay регуляризация L2 будет выполняться автоматически в optimizer.step().

# here is a typical PyTorch epoch
for batch in loader:    
    optimizer.zero_grad()    
    # do train

    loss.backward()
    optimizer.step()

Способ 2: своими руками

Регуляризация заключается в добавлении нормы L1/L2 по отношению к параметры каждого слоя до конечной потери.

Чтобы сделать этот процесс самодельным, нам просто нужно повторить каждый слой, получить параметры и накопить значение нормы до конца и, наконец, сложить его.

Для себя я реализовал функцию, которая принимает model, то есть модель для регуляризации, и тип регуляризации reg_type, который ограничен в пределах L1 или L2 в строке, указывая, что желаемый тип регуляризации и коэффициент регуляризации coef.

import torch.nn as nn
import torch

def regularization(model:nn.Module, reg_type:str, coef:float):
    int_type=int(reg_type[1])
    reg_loss = 0
    for module in model.modules():
        for param in module.parameters():
            reg_loss+=torch.norm(param,int_type)
        
    return reg_loss*coef

Обратите внимание, что функция является чисто пользовательской, я реализовал это с регуляризацией, наложенной на каждый слой модели, но на практике вы можете буквально добавить любое условие if для выбора определенных слоев.