【AI Shift/Kaggle Advent Calendar 2022】Kaggleで学んだBERTをfine-tuningする際のTips④〜Adversarial Training編〜

class FGM():
    def __init__(self, model, criterion, eps):
        self.model = model
        self.backup = {}
        self.eps = eps
        self.criterion = criterion
        
    def __call__(self, inputs, labels):
        self.attack(epsilon=self.eps, emb_name='word_embeddings')
        y_preds = self.model(inputs)
        loss_fgm = self.criterion(y_preds, labels)
        loss_fgm.backward()
        self.restore()

    def attack(self, epsilon=1., emb_name='word_embeddings'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0:
                    r_at = epsilon * param.grad / norm
                    param.data.add_(r_at)
                    
    def restore(self, emb_name='word_embeddings'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                assert name in self.backup
                param.data = self.backup[name]
            self.backup = {}

# 参考: https://www.kaggle.com/code/wht1996/feedback-nn-train/notebook
class AWP:
    def __init__(
        self,
        model,
        optimizer,
        criterion,
        adv_param="weight",
        adv_lr=1,
        adv_eps=0.2,
    ):
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.adv_param = adv_param
        self.adv_lr = adv_lr
        self.adv_eps = adv_eps
        self.backup = {}
        self.backup_eps = {}

    def __call__(self, inputs, labels):
        self._save()
        self._attack_step() 
        y_preds = self.model(inputs)
        adv_loss = self.criterion(y_preds, labels)   
        adv_loss.backward()  
        self._restore()

    def _attack_step(self):
        e = 1e-6
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                norm1 = torch.norm(param.grad)
                norm2 = torch.norm(param.data.detach())
                if norm1 != 0 and not torch.isnan(norm1):
                    r_at = self.adv_lr * param.grad / (norm1 + e) * (norm2 + e)
                    param.data.add_(r_at)
                    param.data = torch.min(
                        torch.max(param.data, self.backup_eps[name][0]), self.backup_eps[name][1]
                    )

    def _save(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                if name not in self.backup:
                    self.backup[name] = param.data.clone()
                    grad_eps = self.adv_eps * param.abs().detach()
                    self.backup_eps[name] = (
                        self.backup[name] - grad_eps,
                        self.backup[name] + grad_eps,
                    )

    def _restore(self,):
        for name, param in self.model.named_parameters():
            if name in self.backup:
                param.data = self.backup[name]
        self.backup = {}
        self.backup_eps = {}

import os
import json
import random
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
from sklearn.metrics import accuracy_score

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset

import transformers
from transformers import AdamW
from transformers import AutoTokenizer, AutoModel, AutoConfig
from transformers import get_cosine_schedule_with_warmup

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class CFG:
    INPUT = "/kaggle/input/nlp-getting-started"
    OUTPUT = "/kaggle/working"
    SEED = 42

    MODEL_NAME = "microsoft/deberta-v3-xsmall"
    TOKENIZER = None
    MAX_LEN = None 
    
    N_EPOCH = 5
    BS = 8
    N_WORKER = 2
    LR = 2e-6
    WEIGHT_DECAY = 0.01
    N_WARMUP = 0
    N_CYCLES = 0.5
    GRAD_NORM = 0.1

    ADV_MODE = "FGM"  # "AWP" or "SiFT" or "baseline" 

TOKENIZER = AutoTokenizer.from_pretrained(CFG.MODEL_NAME)
TOKENIZER.save_pretrained(CFG.OUTPUT+'/tokenizer/')
CFG.TOKENIZER = TOKENIZER
del TOKENIZER

train_df = pd.read_csv(f'{CFG.INPUT}/train.csv')

max_len = 0
for text in train_df['text']:
    tok = CFG.TOKENIZER(text)
    max_len = max(len(tok['input_ids']), max_len)
print('max_len =', max_len). # max_len = 91
CFG.MAX_LEN = max_len

class FB3Dataset(Dataset):
    def __init__(self, df):
        self.texts = df['text'].values
        self.labels = df['target'].values

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, item):
        inputs = CFG.TOKENIZER.encode_plus(
            self.texts[item], 
            return_tensors=None, 
            add_special_tokens=True, 
            max_length=CFG.MAX_LEN,
            padding='max_length',
            truncation=True
        )
        for k, v in inputs.items():
            inputs[k] = torch.tensor(v, dtype=torch.long) 
        label = torch.tensor(self.labels[item], dtype=torch.long)
        return inputs, label

class CustomModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.config = AutoConfig.from_pretrained(CFG.MODEL_NAME, output_hidden_states=True)
        self.model = AutoModel.from_pretrained(CFG.MODEL_NAME, config=self.config)
        self.fc = nn.Linear(self.config.hidden_size, 2)

    def forward(self, inputs):
        outputs = self.model(**inputs)
        last_hidden_states = outputs[0]
        feature = last_hidden_states[:, 0, :]
        output = self.fc(feature)
        
        return output

def train_step(model, train_loader, optimizer, scheduler, criterion, adv):
    losses = []
    model.train()
    for i, (inputs, labels) in tqdm(enumerate(train_loader), total=len(train_loader)):
        optimizer.zero_grad()
        for k in inputs.keys():
            inputs[k] = inputs[k].to(device)
        labels = labels.to(device)
        y_pred = model(inputs)
        loss = criterion(y_pred, labels)
 
        # Adversarial Training
        if CFG.ADV_MODE in ["FGM", "AWP"]:
            loss.backward()
            adv(inputs, labels)  # FGM or AWP
        elif CFG.ADV_MODE == "SiFT":
            def logits_fn(m, *wargs, **kwargs):
                logits = m(inputs)
                return logits
            loss += adv.loss(y_pred, logits_fn, inputs=inputs)
            loss.backward()
        else:
            loss.backward()

        optimizer.step()
        scheduler.step()
        losses.append(loss.cpu().item())
    return losses

def eval_step(model, valid_loader, criterion):
    losses, predicts = [], []
    model.eval()
    for i, (inputs, labels) in tqdm(enumerate(valid_loader), total=len(valid_loader)):
        for k in inputs.keys():
            inputs[k] = inputs[k].to(device)
        labels = labels.to(device)
        with torch.no_grad():
            y_pred = model(inputs)
        loss = criterion(y_pred, labels)
        losses.append(loss.cpu().item())
        predicts.append(y_pred.argmax(1).int().cpu().numpy())
        
    return losses, np.hstack(predicts)

def seed_everything(seed=42):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def make_submission(model):
    model.eval()
    sub_df = pd.read_csv(f'{CFG.INPUT}/sample_submission.csv')
    test_df = pd.read_csv(f'{CFG.INPUT}/test.csv')
    predicts = []
    for text in tqdm(test_df['text']):
        inputs = CFG.TOKENIZER.encode_plus(
            text, 
            return_tensors=None, 
            add_special_tokens=True, 
            max_length=CFG.MAX_LEN,
            padding='max_length',
            truncation=True
        )
        for k, v in inputs.items():
            inputs[k] = torch.tensor([v], dtype=torch.long).to(device)
        with torch.no_grad():
            pred = model(inputs)
        predicts.append(pred.cpu())
    predicts = torch.vstack(predicts).argmax(1).numpy()
    sub_df['target'] = predicts
    sub_df.to_csv('submission.csv', index=None)

def main():
    seed_everything()

    _train_df = train_df.sample(frac=1, random_state=CFG.SEED)
    tes_df = _train_df.iloc[:1500]
    val_df = _train_df.iloc[1500:3000]
    trn_df = _train_df.iloc[3000:]

    y_true_valid = val_df['target'].values
    y_true_test = tes_df['target'].values

    train_dataset = FB3Dataset(trn_df)
    valid_dataset = FB3Dataset(val_df)
    test_dataset = FB3Dataset(tes_df)
    train_loader = DataLoader(train_dataset,
                              batch_size=CFG.BS,
                              shuffle=True,
                              num_workers=CFG.N_WORKER, pin_memory=True, drop_last=True)
    valid_loader = DataLoader(valid_dataset,
                              batch_size=CFG.BS,
                              shuffle=False,
                              num_workers=CFG.N_WORKER, pin_memory=True, drop_last=False)
    test_loader = DataLoader(test_dataset,
                             batch_size=CFG.BS,
                             shuffle=False,
                            num_workers=CFG.N_WORKER, pin_memory=True, drop_last=False)
    num_train_steps = int(len(train_dataset) / CFG.BS * CFG.N_EPOCH)

    model = CustomModel()
    model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = AdamW(model.parameters(), lr=CFG.LR)
    scheduler = get_cosine_schedule_with_warmup(
        optimizer, num_warmup_steps=CFG.N_WARMUP,num_training_steps=num_train_steps, num_cycles=CFG.N_CYCLES
    )
    
    # Adversarialモデルの定義
    if CFG.ADV_MODE == "FGM":
        adv = FGM(model, criterion, 0.1)
    elif CFG.ADV_MODE == "AWP":
        adv = AWP(model,
                  optimizer,
                  criterion,
                  adv_lr=1e-4,
                  adv_eps=1e-4,
        )
    elif CFG.ADV_MODE == "SiFT":
        adv_modules = hook_sift_layer(model, hidden_size=model.config.hidden_size)
        adv = AdversarialLearner(model, adv_modules)
    else:
        adv = None

    best_score = float('-inf')
    for epoch in range(CFG.N_EPOCH):
        print(f"Epoch: {epoch}")
        train_losses = train_step(model, train_loader, optimizer, scheduler, criterion, adv)
        valid_losses, y_pred_valid = eval_step(model, valid_loader, criterion)
        acc = accuracy_score(y_true_valid, y_pred_valid)
        print(f"  valid accuracy={acc}")
        if best_score < acc:
            best_score = acc
            torch.save(model.state_dict(),
                       f"{CFG.OUTPUT}/{CFG.MODEL_NAME.replace('/', '-')}_best.pth")

    state = torch.load(f"{CFG.OUTPUT}/{CFG.MODEL_NAME.replace('/', '-')}_best.pth")
    model.load_state_dict(state)
    _, y_pred_test = eval_step(model, test_loader, criterion)
    acc = accuracy_score(y_true_test, y_pred_test)
    print(f"test score = {acc}")
    make_submission(model)