Upload folder using huggingface_hub

12665fd verified 3 months ago

12.4 kB

"""
Baseline Transformer from scratch.
- Tokenizer SentencePiece BPE construit depuis les données FR+SRR (vocab=8000)
- Architecture Transformer standard (PyTorch nn.Transformer)
- Aucun poids pré-entraîné, aucun vocabulaire externe
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

import json, math
import torch
import torch.nn as nn
import evaluate
import numpy as np
import lightning as L
import sentencepiece as spm
from torch.utils.data import Dataset, DataLoader
from config import Config

TOKENIZER_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tokenizer/spm.model")
PAD_ID, BOS_ID, EOS_ID = 0, 2, 3


# ─── Tokenizer wrapper ────────────────────────────────────────────────────────

class ScratchTokenizer:
    def __init__(self, model_path: str = TOKENIZER_PATH):
        self.sp = spm.SentencePieceProcessor()
        self.sp.load(model_path)
        self.pad_token_id = PAD_ID
        self.bos_token_id = BOS_ID
        self.eos_token_id = EOS_ID
        self.vocab_size   = self.sp.get_piece_size()

    def encode(self, text: str, max_length: int = 128) -> list[int]:
        ids = self.sp.encode(text, out_type=int)
        ids = ids[: max_length - 1] + [EOS_ID]
        return ids

    def decode(self, ids: list[int]) -> str:
        ids = [i for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID)]
        return self.sp.decode(ids)

    def pad(self, sequences: list[list[int]], max_length: int) -> tuple:
        padded = [s + [PAD_ID] * (max_length - len(s)) for s in sequences]
        mask   = [[1] * len(s) + [0] * (max_length - len(s)) for s in sequences]
        return padded, mask


# ─── Dataset ──────────────────────────────────────────────────────────────────

class ScratchDataset(Dataset):
    def __init__(self, path: str, tokenizer: ScratchTokenizer, cfg: Config):
        with open(path, encoding="utf-8") as f:
            self.data = json.load(f)
        self.tokenizer = tokenizer
        self.cfg = cfg

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

    def __getitem__(self, idx):
        item = self.data[idx]
        src_ids = self.tokenizer.encode(item["francais"], self.cfg.MAX_LENGTH)
        tgt_ids = [BOS_ID] + self.tokenizer.encode(item["serere"], self.cfg.MAX_LENGTH)

        src_ids = (src_ids + [PAD_ID] * self.cfg.MAX_LENGTH)[: self.cfg.MAX_LENGTH]
        tgt_ids = (tgt_ids + [PAD_ID] * (self.cfg.MAX_LENGTH + 1))[: self.cfg.MAX_LENGTH + 1]

        src_ids = torch.tensor(src_ids, dtype=torch.long)
        tgt_ids = torch.tensor(tgt_ids, dtype=torch.long)

        src_mask = (src_ids != PAD_ID).long()
        labels   = tgt_ids[1:].clone()
        labels[labels == PAD_ID] = -100
        decoder_input = tgt_ids[:-1]

        return {
            "input_ids":      src_ids,
            "attention_mask": src_mask,
            "decoder_input":  decoder_input,
            "labels":         labels,
        }


class ScratchDataModule(L.LightningDataModule):
    def __init__(self, tokenizer: ScratchTokenizer, cfg: Config):
        super().__init__()
        self.tokenizer = tokenizer
        self.cfg = cfg

    def setup(self, stage=None):
        self.train_ds = ScratchDataset(
            os.path.join(self.cfg.DATA_DIR, "train.json"), self.tokenizer, self.cfg
        )
        self.val_ds = ScratchDataset(
            os.path.join(self.cfg.DATA_DIR, "val.json"), self.tokenizer, self.cfg
        )

    def train_dataloader(self):
        return DataLoader(self.train_ds, batch_size=self.cfg.BATCH_SIZE, shuffle=True,  num_workers=4, pin_memory=True)

    def val_dataloader(self):
        return DataLoader(self.val_ds,   batch_size=self.cfg.BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)


# ─── Positional Encoding ──────────────────────────────────────────────────────

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(max_len).unsqueeze(1)
        div = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        return self.dropout(x + self.pe[:, : x.size(1)])


# ─── Seq2Seq Transformer ──────────────────────────────────────────────────────

class Seq2SeqTransformer(nn.Module):
    def __init__(self, vocab_size: int, d_model=256, nhead=8,
                 num_enc=3, num_dec=3, ffn=512, dropout=0.1, max_len=256):
        super().__init__()
        self.src_emb = nn.Embedding(vocab_size, d_model, padding_idx=PAD_ID)
        self.tgt_emb = nn.Embedding(vocab_size, d_model, padding_idx=PAD_ID)
        self.pos_enc = PositionalEncoding(d_model, max_len, dropout)
        self.transformer = nn.Transformer(
            d_model=d_model, nhead=nhead,
            num_encoder_layers=num_enc, num_decoder_layers=num_dec,
            dim_feedforward=ffn, dropout=dropout, batch_first=True,
        )
        self.proj = nn.Linear(d_model, vocab_size)
        self._init_weights()

    def _init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, src, tgt, src_key_padding_mask=None):
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1), device=src.device)
        tgt_pad_mask = (tgt == PAD_ID)
        src_emb = self.pos_enc(self.src_emb(src) * math.sqrt(self.src_emb.embedding_dim))
        tgt_emb = self.pos_enc(self.tgt_emb(tgt) * math.sqrt(self.tgt_emb.embedding_dim))
        out = self.transformer(
            src_emb, tgt_emb,
            tgt_mask=tgt_mask,
            src_key_padding_mask=(src_key_padding_mask == 0) if src_key_padding_mask is not None else None,
            tgt_key_padding_mask=tgt_pad_mask,
            memory_key_padding_mask=(src_key_padding_mask == 0) if src_key_padding_mask is not None else None,
        )
        return self.proj(out)

    @torch.no_grad()
    def generate(self, src, attention_mask=None, max_new_tokens=64, num_beams=1):
        if num_beams != 1:
            raise NotImplementedError(
                "Le Transformer scratch utilise le décodage glouton; num_beams doit valoir 1."
            )
        B = src.size(0)
        memory_pad = (attention_mask == 0) if attention_mask is not None else None
        src_emb = self.pos_enc(self.src_emb(src) * math.sqrt(self.src_emb.embedding_dim))
        memory  = self.transformer.encoder(
            src_emb,
            src_key_padding_mask=memory_pad,
        )
        ys = torch.full((B, 1), BOS_ID, dtype=torch.long, device=src.device)
        finished = torch.zeros(B, dtype=torch.bool, device=src.device)
        for _ in range(max_new_tokens):
            tgt_emb = self.pos_enc(self.tgt_emb(ys) * math.sqrt(self.tgt_emb.embedding_dim))
            tgt_mask = nn.Transformer.generate_square_subsequent_mask(ys.size(1), device=src.device)
            out  = self.transformer.decoder(tgt_emb, memory, tgt_mask=tgt_mask,
                                            memory_key_padding_mask=memory_pad)
            logits   = self.proj(out[:, -1, :])
            next_tok = logits.argmax(-1, keepdim=True)
            # Après le premier EOS d'une phrase, produire uniquement PAD sans
            # modifier les autres phrases encore actives du batch.
            next_tok = torch.where(
                finished.unsqueeze(1), torch.full_like(next_tok, PAD_ID), next_tok
            )
            ys = torch.cat([ys, next_tok], dim=1)
            finished |= next_tok.squeeze(1).eq(EOS_ID)
            if finished.all():
                break
        return ys


# ─── Lightning Module ─────────────────────────────────────────────────────────

class ScratchTransformer(L.LightningModule):
    def __init__(self, cfg: Config):
        super().__init__()
        self.save_hyperparameters()
        self.cfg = cfg
        if not os.path.exists(TOKENIZER_PATH):
            raise FileNotFoundError(
                f"Tokenizer introuvable : {TOKENIZER_PATH}\n"
                "Lancez d'abord : python baselines/build_tokenizer.py"
            )
        self.tokenizer = ScratchTokenizer(TOKENIZER_PATH)
        self.model = Seq2SeqTransformer(
            self.tokenizer.vocab_size,
            d_model=cfg.D_MODEL,
            nhead=cfg.N_HEADS,
            num_enc=cfg.N_ENCODER_LAYERS,
            num_dec=cfg.N_DECODER_LAYERS,
            ffn=cfg.D_FF,
            dropout=cfg.DROPOUT,
            max_len=cfg.MAX_LENGTH + 2,
        )
        self.criterion = nn.CrossEntropyLoss(ignore_index=-100)

        self.metric_bleu      = evaluate.load("sacrebleu")
        self.metric_rouge     = evaluate.load("rouge")
        self.metric_bertscore = evaluate.load("bertscore")
        self.validation_step_outputs = []

    def forward(self, src, tgt_in, src_mask=None):
        return self.model(src, tgt_in, src_key_padding_mask=src_mask)

    def training_step(self, batch, batch_idx):
        logits = self(batch["input_ids"], batch["decoder_input"], batch["attention_mask"])
        loss   = self.criterion(logits.reshape(-1, self.tokenizer.vocab_size), batch["labels"].reshape(-1))
        self.log("train_loss", loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        logits = self(batch["input_ids"], batch["decoder_input"], batch["attention_mask"])
        loss   = self.criterion(logits.reshape(-1, self.tokenizer.vocab_size), batch["labels"].reshape(-1))
        self.log("val_loss", loss, prog_bar=True, sync_dist=True)

        generated = self.model.generate(
            batch["input_ids"], batch["attention_mask"], max_new_tokens=self.cfg.MAX_LENGTH
        )
        preds = [self.tokenizer.decode(g.tolist()) for g in generated]
        refs  = [self.tokenizer.decode(
            torch.where(l != -100, l, torch.tensor(PAD_ID)).tolist()
        ) for l in batch["labels"]]
        self.validation_step_outputs.append({"preds": preds, "targets": refs})
        return loss

    def on_validation_epoch_end(self):
        all_preds   = [p for x in self.validation_step_outputs for p in x["preds"]]
        all_targets = [t for x in self.validation_step_outputs for t in x["targets"]]
        if not all_preds:
            return
        bleu  = self.metric_bleu.compute(predictions=all_preds, references=all_targets)
        rouge = self.metric_rouge.compute(predictions=all_preds, references=all_targets)
        bert  = self.metric_bertscore.compute(
            predictions=all_preds, references=all_targets,
            model_type=self.cfg.BERTSCORE_MODEL, lang=self.cfg.BERTSCORE_LANG,
        )
        # sync_dist=True: voir models/model_lora.py pour l'explication (course
        # DDP sur val_bleu -> checkpoints incohérents entre ranks).
        self.log("val_bleu",         bleu["score"],        prog_bar=True, sync_dist=True)
        self.log("val_rouge1",       rouge["rouge1"],       sync_dist=True)
        self.log("val_rougeL",       rouge["rougeL"],       sync_dist=True)
        self.log("val_bertscore_f1", np.mean(bert["f1"]),   sync_dist=True)
        self.validation_step_outputs.clear()

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.cfg.LR, weight_decay=self.cfg.WEIGHT_DECAY)
        if self.cfg.WARMUP_STEPS <= 0:
            return optimizer
        scheduler = torch.optim.lr_scheduler.LinearLR(
            optimizer, start_factor=0.1, end_factor=1.0, total_iters=self.cfg.WARMUP_STEPS
        )
        return [optimizer], [{"scheduler": scheduler, "interval": "step"}]

	"""
	Baseline Transformer from scratch.
	- Tokenizer SentencePiece BPE construit depuis les données FR+SRR (vocab=8000)
	- Architecture Transformer standard (PyTorch nn.Transformer)
	- Aucun poids pré-entraîné, aucun vocabulaire externe
	"""
	import sys, os
	sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

	import json, math
	import torch
	import torch.nn as nn
	import evaluate
	import numpy as np
	import lightning as L
	import sentencepiece as spm
	from torch.utils.data import Dataset, DataLoader
	from config import Config

	TOKENIZER_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tokenizer/spm.model")
	PAD_ID, BOS_ID, EOS_ID = 0, 2, 3


	# ─── Tokenizer wrapper ────────────────────────────────────────────────────────

	class ScratchTokenizer:
	def __init__(self, model_path: str = TOKENIZER_PATH):
	self.sp = spm.SentencePieceProcessor()
	self.sp.load(model_path)
	self.pad_token_id = PAD_ID
	self.bos_token_id = BOS_ID
	self.eos_token_id = EOS_ID
	self.vocab_size = self.sp.get_piece_size()

	def encode(self, text: str, max_length: int = 128) -> list[int]:
	ids = self.sp.encode(text, out_type=int)
	ids = ids[: max_length - 1] + [EOS_ID]
	return ids

	def decode(self, ids: list[int]) -> str:
	ids = [i for i in ids if i not in (PAD_ID, BOS_ID, EOS_ID)]
	return self.sp.decode(ids)

	def pad(self, sequences: list[list[int]], max_length: int) -> tuple:
	padded = [s + [PAD_ID] * (max_length - len(s)) for s in sequences]
	mask = [[1] * len(s) + [0] * (max_length - len(s)) for s in sequences]
	return padded, mask


	# ─── Dataset ──────────────────────────────────────────────────────────────────

	class ScratchDataset(Dataset):
	def __init__(self, path: str, tokenizer: ScratchTokenizer, cfg: Config):
	with open(path, encoding="utf-8") as f:
	self.data = json.load(f)
	self.tokenizer = tokenizer
	self.cfg = cfg

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

	def __getitem__(self, idx):
	item = self.data[idx]
	src_ids = self.tokenizer.encode(item["francais"], self.cfg.MAX_LENGTH)
	tgt_ids = [BOS_ID] + self.tokenizer.encode(item["serere"], self.cfg.MAX_LENGTH)

	src_ids = (src_ids + [PAD_ID] * self.cfg.MAX_LENGTH)[: self.cfg.MAX_LENGTH]
	tgt_ids = (tgt_ids + [PAD_ID] * (self.cfg.MAX_LENGTH + 1))[: self.cfg.MAX_LENGTH + 1]

	src_ids = torch.tensor(src_ids, dtype=torch.long)
	tgt_ids = torch.tensor(tgt_ids, dtype=torch.long)

	src_mask = (src_ids != PAD_ID).long()
	labels = tgt_ids[1:].clone()
	labels[labels == PAD_ID] = -100
	decoder_input = tgt_ids[:-1]

	return {
	"input_ids": src_ids,
	"attention_mask": src_mask,
	"decoder_input": decoder_input,
	"labels": labels,
	}


	class ScratchDataModule(L.LightningDataModule):
	def __init__(self, tokenizer: ScratchTokenizer, cfg: Config):
	super().__init__()
	self.tokenizer = tokenizer
	self.cfg = cfg

	def setup(self, stage=None):
	self.train_ds = ScratchDataset(
	os.path.join(self.cfg.DATA_DIR, "train.json"), self.tokenizer, self.cfg
	)
	self.val_ds = ScratchDataset(
	os.path.join(self.cfg.DATA_DIR, "val.json"), self.tokenizer, self.cfg
	)

	def train_dataloader(self):
	return DataLoader(self.train_ds, batch_size=self.cfg.BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True)

	def val_dataloader(self):
	return DataLoader(self.val_ds, batch_size=self.cfg.BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True)


	# ─── Positional Encoding ──────────────────────────────────────────────────────

	class PositionalEncoding(nn.Module):
	def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
	super().__init__()
	self.dropout = nn.Dropout(dropout)
	pe = torch.zeros(max_len, d_model)
	pos = torch.arange(max_len).unsqueeze(1)
	div = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
	pe[:, 0::2] = torch.sin(pos * div)
	pe[:, 1::2] = torch.cos(pos * div)
	self.register_buffer("pe", pe.unsqueeze(0))

	def forward(self, x):
	return self.dropout(x + self.pe[:, : x.size(1)])


	# ─── Seq2Seq Transformer ──────────────────────────────────────────────────────

	class Seq2SeqTransformer(nn.Module):
	def __init__(self, vocab_size: int, d_model=256, nhead=8,
	num_enc=3, num_dec=3, ffn=512, dropout=0.1, max_len=256):
	super().__init__()
	self.src_emb = nn.Embedding(vocab_size, d_model, padding_idx=PAD_ID)
	self.tgt_emb = nn.Embedding(vocab_size, d_model, padding_idx=PAD_ID)
	self.pos_enc = PositionalEncoding(d_model, max_len, dropout)
	self.transformer = nn.Transformer(
	d_model=d_model, nhead=nhead,
	num_encoder_layers=num_enc, num_decoder_layers=num_dec,
	dim_feedforward=ffn, dropout=dropout, batch_first=True,
	)
	self.proj = nn.Linear(d_model, vocab_size)
	self._init_weights()

	def _init_weights(self):
	for p in self.parameters():
	if p.dim() > 1:
	nn.init.xavier_uniform_(p)

	def forward(self, src, tgt, src_key_padding_mask=None):
	tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1), device=src.device)
	tgt_pad_mask = (tgt == PAD_ID)
	src_emb = self.pos_enc(self.src_emb(src) * math.sqrt(self.src_emb.embedding_dim))
	tgt_emb = self.pos_enc(self.tgt_emb(tgt) * math.sqrt(self.tgt_emb.embedding_dim))
	out = self.transformer(
	src_emb, tgt_emb,
	tgt_mask=tgt_mask,
	src_key_padding_mask=(src_key_padding_mask == 0) if src_key_padding_mask is not None else None,
	tgt_key_padding_mask=tgt_pad_mask,
	memory_key_padding_mask=(src_key_padding_mask == 0) if src_key_padding_mask is not None else None,
	)
	return self.proj(out)

	@torch.no_grad()
	def generate(self, src, attention_mask=None, max_new_tokens=64, num_beams=1):
	if num_beams != 1:
	raise NotImplementedError(
	"Le Transformer scratch utilise le décodage glouton; num_beams doit valoir 1."
	)
	B = src.size(0)
	memory_pad = (attention_mask == 0) if attention_mask is not None else None
	src_emb = self.pos_enc(self.src_emb(src) * math.sqrt(self.src_emb.embedding_dim))
	memory = self.transformer.encoder(
	src_emb,
	src_key_padding_mask=memory_pad,
	)
	ys = torch.full((B, 1), BOS_ID, dtype=torch.long, device=src.device)
	finished = torch.zeros(B, dtype=torch.bool, device=src.device)
	for _ in range(max_new_tokens):
	tgt_emb = self.pos_enc(self.tgt_emb(ys) * math.sqrt(self.tgt_emb.embedding_dim))
	tgt_mask = nn.Transformer.generate_square_subsequent_mask(ys.size(1), device=src.device)
	out = self.transformer.decoder(tgt_emb, memory, tgt_mask=tgt_mask,
	memory_key_padding_mask=memory_pad)
	logits = self.proj(out[:, -1, :])
	next_tok = logits.argmax(-1, keepdim=True)
	# Après le premier EOS d'une phrase, produire uniquement PAD sans
	# modifier les autres phrases encore actives du batch.
	next_tok = torch.where(
	finished.unsqueeze(1), torch.full_like(next_tok, PAD_ID), next_tok
	)
	ys = torch.cat([ys, next_tok], dim=1)
	finished \|= next_tok.squeeze(1).eq(EOS_ID)
	if finished.all():
	break
	return ys


	# ─── Lightning Module ─────────────────────────────────────────────────────────

	class ScratchTransformer(L.LightningModule):
	def __init__(self, cfg: Config):
	super().__init__()
	self.save_hyperparameters()
	self.cfg = cfg
	if not os.path.exists(TOKENIZER_PATH):
	raise FileNotFoundError(
	f"Tokenizer introuvable : {TOKENIZER_PATH}\n"
	"Lancez d'abord : python baselines/build_tokenizer.py"
	)
	self.tokenizer = ScratchTokenizer(TOKENIZER_PATH)
	self.model = Seq2SeqTransformer(
	self.tokenizer.vocab_size,
	d_model=cfg.D_MODEL,
	nhead=cfg.N_HEADS,
	num_enc=cfg.N_ENCODER_LAYERS,
	num_dec=cfg.N_DECODER_LAYERS,
	ffn=cfg.D_FF,
	dropout=cfg.DROPOUT,
	max_len=cfg.MAX_LENGTH + 2,
	)
	self.criterion = nn.CrossEntropyLoss(ignore_index=-100)

	self.metric_bleu = evaluate.load("sacrebleu")
	self.metric_rouge = evaluate.load("rouge")
	self.metric_bertscore = evaluate.load("bertscore")
	self.validation_step_outputs = []

	def forward(self, src, tgt_in, src_mask=None):
	return self.model(src, tgt_in, src_key_padding_mask=src_mask)

	def training_step(self, batch, batch_idx):
	logits = self(batch["input_ids"], batch["decoder_input"], batch["attention_mask"])
	loss = self.criterion(logits.reshape(-1, self.tokenizer.vocab_size), batch["labels"].reshape(-1))
	self.log("train_loss", loss, prog_bar=True)
	return loss

	def validation_step(self, batch, batch_idx):
	logits = self(batch["input_ids"], batch["decoder_input"], batch["attention_mask"])
	loss = self.criterion(logits.reshape(-1, self.tokenizer.vocab_size), batch["labels"].reshape(-1))
	self.log("val_loss", loss, prog_bar=True, sync_dist=True)

	generated = self.model.generate(
	batch["input_ids"], batch["attention_mask"], max_new_tokens=self.cfg.MAX_LENGTH
	)
	preds = [self.tokenizer.decode(g.tolist()) for g in generated]
	refs = [self.tokenizer.decode(
	torch.where(l != -100, l, torch.tensor(PAD_ID)).tolist()
	) for l in batch["labels"]]
	self.validation_step_outputs.append({"preds": preds, "targets": refs})
	return loss

	def on_validation_epoch_end(self):
	all_preds = [p for x in self.validation_step_outputs for p in x["preds"]]
	all_targets = [t for x in self.validation_step_outputs for t in x["targets"]]
	if not all_preds:
	return
	bleu = self.metric_bleu.compute(predictions=all_preds, references=all_targets)
	rouge = self.metric_rouge.compute(predictions=all_preds, references=all_targets)
	bert = self.metric_bertscore.compute(
	predictions=all_preds, references=all_targets,
	model_type=self.cfg.BERTSCORE_MODEL, lang=self.cfg.BERTSCORE_LANG,
	)
	# sync_dist=True: voir models/model_lora.py pour l'explication (course
	# DDP sur val_bleu -> checkpoints incohérents entre ranks).
	self.log("val_bleu", bleu["score"], prog_bar=True, sync_dist=True)
	self.log("val_rouge1", rouge["rouge1"], sync_dist=True)
	self.log("val_rougeL", rouge["rougeL"], sync_dist=True)
	self.log("val_bertscore_f1", np.mean(bert["f1"]), sync_dist=True)
	self.validation_step_outputs.clear()

	def configure_optimizers(self):
	optimizer = torch.optim.AdamW(self.parameters(), lr=self.cfg.LR, weight_decay=self.cfg.WEIGHT_DECAY)
	if self.cfg.WARMUP_STEPS <= 0:
	return optimizer
	scheduler = torch.optim.lr_scheduler.LinearLR(
	optimizer, start_factor=0.1, end_factor=1.0, total_iters=self.cfg.WARMUP_STEPS
	)
	return [optimizer], [{"scheduler": scheduler, "interval": "step"}]