Kaggleで学んだBERTをfine-tuningする際のTips③〜過学習抑制編〜

class CommonLitModel(nn.Module):
    
    def __init__(self):
        super(CommonLitModel, self).__init__()
        self.config = AutoConfig.from_pretrained(MODEL_NAME)
        self.bert = AutoModel.from_pretrained(
            MODEL_NAME,
            config=self.config
        )
        self.dropouts = nn.ModuleList([nn.Dropout(0.2) for _ in range(N_MSD)])
        self.regressor = nn.Linear(self.config.hidden_size, 1)

    def forward(self, input_ids, attention_mask, token_type_ids):
        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        sequence_output = outputs['last_hidden_state'][:, 0]  # baseline
        logits = sum([self.regressor(dropout(sequence_output)) for dropout in self.dropouts])/N_MSD

        return logits

import math
from torch.autograd.function import InplaceFunction
from torch.nn import Parameter
import torch.nn.init as init

class Mixout(InplaceFunction):
    @staticmethod
    def _make_noise(input):
        return input.new().resize_as_(input)

    @classmethod
    def forward(cls, ctx, input, target=None, p=0.0, training=False, inplace=False):
        if p < 0 or p > 1:
            raise ValueError("A mix probability of mixout has to be between 0 and 1," " but got {}".format(p))
        if target is not None and input.size() != target.size():
            raise ValueError(
                "A target tensor size must match with a input tensor size {},"
                " but got {}".format(input.size(), target.size())
            )
        ctx.p = p
        ctx.training = training

        if ctx.p == 0 or not ctx.training:
            return input

        if target is None:
            target = cls._make_noise(input)
            target.fill_(0)
        target = target.to(input.device)

        if inplace:
            ctx.mark_dirty(input)
            output = input
        else:
            output = input.clone()

        ctx.noise = cls._make_noise(input)
        if len(ctx.noise.size()) == 1:
            ctx.noise.bernoulli_(1 - ctx.p)
        else:
            ctx.noise[0].bernoulli_(1 - ctx.p)
            ctx.noise = ctx.noise[0].repeat(input.size()[0], 1)
        ctx.noise.expand_as(input)

        if ctx.p == 1:
            output = target
        else:
            output = ((1 - ctx.noise) * target + ctx.noise * output - ctx.p * target) / (1 - ctx.p)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        if ctx.p > 0 and ctx.training:
            return grad_output * ctx.noise, None, None, None, None
        else:
            return grad_output, None, None, None, None


def mixout(input, target=None, p=0.0, training=False, inplace=False):
    return Mixout.apply(input, target, p, training, inplace)


class MixLinear(torch.nn.Module):
    __constants__ = ["bias", "in_features", "out_features"]
    def __init__(self, in_features, out_features, bias=True, target=None, p=0.0):
        super(MixLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.Tensor(out_features, in_features))
        if bias:
            self.bias = Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter("bias", None)
        self.reset_parameters()
        self.target = target
        self.p = p

    def reset_parameters(self):
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)

    def forward(self, input):
        return F.linear(input, mixout(self.weight, self.target, self.p, self.training), self.bias)

    def extra_repr(self):
        type = "drop" if self.target is None else "mix"
        return "{}={}, in_features={}, out_features={}, bias={}".format(
            type + "out", self.p, self.in_features, self.out_features, self.bias is not None
        )

def replace_mixout(model):
    for sup_module in model.modules():
        for name, module in sup_module.named_children():
            if isinstance(module, nn.Dropout):
                module.p = 0.0
            if isinstance(module, nn.Linear):
                target_state_dict = module.state_dict()
                bias = True if module.bias is not None else False
                new_module = MixLinear(
                    module.in_features, module.out_features, bias, target_state_dict["weight"], MIXOUT
                )
                new_module.load_state_dict(target_state_dict)
                setattr(sup_module, name, new_module)
    return model

model = CommonLitModel()
+model = replace_mixout(model)
model.to(device)

def compute_kl_loss(p, q, pad_mask=None):

    p_loss = F.kl_div(F.log_softmax(p, dim=-1), F.softmax(q, dim=-1), reduction='none')
    q_loss = F.kl_div(F.log_softmax(q, dim=-1), F.softmax(p, dim=-1), reduction='none')

    # pad_mask is for seq-level tasks
    if pad_mask is not None:
        p_loss.masked_fill_(pad_mask, 0.)
        q_loss.masked_fill_(pad_mask, 0.)

    # You can choose whether to use function "sum" and "mean" depending on your task
    p_loss = p_loss.sum()
    q_loss = q_loss.sum()

    loss = (p_loss + q_loss) / 2
    return loss

def loss_fn(logits1, logits2, label):
    # RMSE loss
    rmse1 = torch.sqrt(nn.MSELoss(reduction='mean')(logits1[:, 0], label))
    rmse2 = torch.sqrt(nn.MSELoss(reduction='mean')(logits2[:, 0], label))

    # R-Drop
    kl_loss = compute_kl_loss(rmse1, rmse2)

    total_loss = (rmse1 + rmse2) / 2 + RDROP_ALPHA * kl_loss
    return total_loss

logits1 = model(
    d["input_ids"].to(device),
    d["attention_mask"].to(device),
    d["token_type_ids"].to(device)
)
logits2 = model(
    d["input_ids"].to(device),
    d["attention_mask"].to(device),
    d["token_type_ids"].to(device)
)
loss= loss_fn(logits1, logits2, d["label"].float().to(device))

class CommonLitModel(nn.Module):
    
    def __init__(self):
        super(CommonLitModel, self).__init__()
        self.config = AutoConfig.from_pretrained(MODEL_NAME)
+        self.config.attention_probs_dropout_prob = 0.0
+        self.config.hidden_dropout_prob = 0.0
        self.bert = AutoModel.from_pretrained(
            MODEL_NAME,
            config=self.config
        )