|
"""Training code for the detector model""" |
|
|
|
import argparse |
|
import os |
|
import subprocess |
|
import sys |
|
from itertools import count |
|
from multiprocessing import Process |
|
|
|
import torch |
|
import torch.distributed as dist |
|
from torch import nn |
|
from torch.nn.parallel import DistributedDataParallel |
|
from torch.optim import Adam |
|
from torch.utils.data import DataLoader, DistributedSampler, RandomSampler |
|
from tqdm import tqdm |
|
from transformers import * |
|
|
|
from .dataset import Corpus, EncodedDataset |
|
from .download import download |
|
from .utils import summary, distributed |
|
|
|
|
|
def setup_distributed(port=29500): |
|
if not dist.is_available() or not torch.cuda.is_available() or torch.cuda.device_count() <= 1: |
|
return 0, 1 |
|
|
|
if 'MPIR_CVAR_CH3_INTERFACE_HOSTNAME' in os.environ: |
|
from mpi4py import MPI |
|
mpi_rank = MPI.COMM_WORLD.Get_rank() |
|
mpi_size = MPI.COMM_WORLD.Get_size() |
|
|
|
os.environ["MASTER_ADDR"] = '127.0.0.1' |
|
os.environ["MASTER_PORT"] = str(port) |
|
|
|
dist.init_process_group(backend="nccl", world_size=mpi_size, rank=mpi_rank) |
|
return mpi_rank, mpi_size |
|
|
|
dist.init_process_group(backend="nccl", init_method="env://") |
|
return dist.get_rank(), dist.get_world_size() |
|
|
|
|
|
def load_datasets(data_dir, real_dataset, fake_dataset, tokenizer, batch_size, |
|
max_sequence_length, random_sequence_length, epoch_size=None, token_dropout=None, seed=None): |
|
if fake_dataset == 'TWO': |
|
download(real_dataset, 'xl-1542M', 'xl-1542M-nucleus', data_dir=data_dir) |
|
elif fake_dataset == 'THREE': |
|
download(real_dataset, 'xl-1542M', 'xl-1542M-k40', 'xl-1542M-nucleus', data_dir=data_dir) |
|
else: |
|
download(real_dataset, fake_dataset, data_dir=data_dir) |
|
|
|
real_corpus = Corpus(real_dataset, data_dir=data_dir) |
|
|
|
if fake_dataset == "TWO": |
|
real_train, real_valid = real_corpus.train * 2, real_corpus.valid * 2 |
|
fake_corpora = [Corpus(name, data_dir=data_dir) for name in ['xl-1542M', 'xl-1542M-nucleus']] |
|
fake_train = sum([corpus.train for corpus in fake_corpora], []) |
|
fake_valid = sum([corpus.valid for corpus in fake_corpora], []) |
|
elif fake_dataset == "THREE": |
|
real_train, real_valid = real_corpus.train * 3, real_corpus.valid * 3 |
|
fake_corpora = [Corpus(name, data_dir=data_dir) for name in |
|
['xl-1542M', 'xl-1542M-k40', 'xl-1542M-nucleus']] |
|
fake_train = sum([corpus.train for corpus in fake_corpora], []) |
|
fake_valid = sum([corpus.valid for corpus in fake_corpora], []) |
|
else: |
|
fake_corpus = Corpus(fake_dataset, data_dir=data_dir) |
|
|
|
real_train, real_valid = real_corpus.train, real_corpus.valid |
|
fake_train, fake_valid = fake_corpus.train, fake_corpus.valid |
|
|
|
Sampler = DistributedSampler if distributed() and dist.get_world_size() > 1 else RandomSampler |
|
|
|
min_sequence_length = 10 if random_sequence_length else None |
|
train_dataset = EncodedDataset(real_train, fake_train, tokenizer, max_sequence_length, min_sequence_length, |
|
epoch_size, token_dropout, seed) |
|
train_loader = DataLoader(train_dataset, batch_size, sampler=Sampler(train_dataset), num_workers=0) |
|
|
|
validation_dataset = EncodedDataset(real_valid, fake_valid, tokenizer) |
|
validation_loader = DataLoader(validation_dataset, batch_size=1, sampler=Sampler(validation_dataset)) |
|
|
|
return train_loader, validation_loader |
|
|
|
|
|
def accuracy_sum(logits, labels): |
|
if list(logits.shape) == list(labels.shape) + [2]: |
|
|
|
classification = (logits[..., 0] < logits[..., 1]).long().flatten() |
|
else: |
|
classification = (logits > 0).long().flatten() |
|
assert classification.shape == labels.shape |
|
return (classification == labels).float().sum().item() |
|
|
|
|
|
def train(model: nn.Module, optimizer, device: str, loader: DataLoader, desc='Train'): |
|
model.train() |
|
|
|
train_accuracy = 0 |
|
train_epoch_size = 0 |
|
train_loss = 0 |
|
|
|
with tqdm(loader, desc=desc, disable=distributed() and dist.get_rank() > 0) as loop: |
|
for texts, masks, labels in loop: |
|
|
|
texts, masks, labels = texts.to(device), masks.to(device), labels.to(device) |
|
batch_size = texts.shape[0] |
|
|
|
optimizer.zero_grad() |
|
loss, logits = model(texts, attention_mask=masks, labels=labels) |
|
loss.backward() |
|
optimizer.step() |
|
|
|
batch_accuracy = accuracy_sum(logits, labels) |
|
train_accuracy += batch_accuracy |
|
train_epoch_size += batch_size |
|
train_loss += loss.item() * batch_size |
|
|
|
loop.set_postfix(loss=loss.item(), acc=train_accuracy / train_epoch_size) |
|
|
|
return { |
|
"train/accuracy": train_accuracy, |
|
"train/epoch_size": train_epoch_size, |
|
"train/loss": train_loss |
|
} |
|
|
|
|
|
def validate(model: nn.Module, device: str, loader: DataLoader, votes=1, desc='Validation'): |
|
model.eval() |
|
|
|
validation_accuracy = 0 |
|
validation_epoch_size = 0 |
|
validation_loss = 0 |
|
|
|
records = [record for v in range(votes) for record in tqdm(loader, desc=f'Preloading data ... {v}', |
|
disable=dist.is_available() and dist.get_rank() > 0)] |
|
records = [[records[v * len(loader) + i] for v in range(votes)] for i in range(len(loader))] |
|
|
|
with tqdm(records, desc=desc, disable=distributed() and dist.get_rank() > 0) as loop, torch.no_grad(): |
|
for example in loop: |
|
losses = [] |
|
logit_votes = [] |
|
|
|
for texts, masks, labels in example: |
|
texts, masks, labels = texts.to(device), masks.to(device), labels.to(device) |
|
batch_size = texts.shape[0] |
|
|
|
loss, logits = model(texts, attention_mask=masks, labels=labels) |
|
losses.append(loss) |
|
logit_votes.append(logits) |
|
|
|
loss = torch.stack(losses).mean(dim=0) |
|
logits = torch.stack(logit_votes).mean(dim=0) |
|
|
|
batch_accuracy = accuracy_sum(logits, labels) |
|
validation_accuracy += batch_accuracy |
|
validation_epoch_size += batch_size |
|
validation_loss += loss.item() * batch_size |
|
|
|
loop.set_postfix(loss=loss.item(), acc=validation_accuracy / validation_epoch_size) |
|
|
|
return { |
|
"validation/accuracy": validation_accuracy, |
|
"validation/epoch_size": validation_epoch_size, |
|
"validation/loss": validation_loss |
|
} |
|
|
|
|
|
def _all_reduce_dict(d, device): |
|
|
|
output_d = {} |
|
for (key, value) in sorted(d.items()): |
|
tensor_input = torch.tensor([[value]]).to(device) |
|
torch.distributed.all_reduce(tensor_input) |
|
output_d[key] = tensor_input.item() |
|
return output_d |
|
|
|
|
|
def run(max_epochs=None, |
|
device=None, |
|
batch_size=24, |
|
max_sequence_length=128, |
|
random_sequence_length=False, |
|
epoch_size=None, |
|
seed=None, |
|
data_dir='data', |
|
real_dataset='webtext', |
|
fake_dataset='xl-1542M-nucleus', |
|
token_dropout=None, |
|
large=False, |
|
learning_rate=2e-5, |
|
weight_decay=0, |
|
**kwargs): |
|
args = locals() |
|
rank, world_size = setup_distributed() |
|
|
|
if device is None: |
|
device = f'cuda:{rank}' if torch.cuda.is_available() else 'cpu' |
|
|
|
print('rank:', rank, 'world_size:', world_size, 'device:', device) |
|
|
|
import torch.distributed as dist |
|
if distributed() and rank > 0: |
|
dist.barrier() |
|
|
|
model_name = 'roberta-large' if large else 'roberta-base' |
|
tokenization_utils.logger.setLevel('ERROR') |
|
tokenizer = RobertaTokenizer.from_pretrained(model_name) |
|
model = RobertaForSequenceClassification.from_pretrained(model_name).to(device) |
|
|
|
if rank == 0: |
|
summary(model) |
|
if distributed(): |
|
dist.barrier() |
|
|
|
if world_size > 1: |
|
model = DistributedDataParallel(model, [rank], output_device=rank, find_unused_parameters=True) |
|
|
|
train_loader, validation_loader = load_datasets(data_dir, real_dataset, fake_dataset, tokenizer, batch_size, |
|
max_sequence_length, random_sequence_length, epoch_size, |
|
token_dropout, seed) |
|
|
|
optimizer = Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay) |
|
epoch_loop = count(1) if max_epochs is None else range(1, max_epochs + 1) |
|
|
|
logdir = os.environ.get("OPENAI_LOGDIR", "logs") |
|
os.makedirs(logdir, exist_ok=True) |
|
|
|
from torch.utils.tensorboard import SummaryWriter |
|
writer = SummaryWriter(logdir) if rank == 0 else None |
|
best_validation_accuracy = 0 |
|
|
|
for epoch in epoch_loop: |
|
if world_size > 1: |
|
train_loader.sampler.set_epoch(epoch) |
|
validation_loader.sampler.set_epoch(epoch) |
|
|
|
train_metrics = train(model, optimizer, device, train_loader, f'Epoch {epoch}') |
|
validation_metrics = validate(model, device, validation_loader) |
|
|
|
combined_metrics = _all_reduce_dict({**validation_metrics, **train_metrics}, device) |
|
|
|
combined_metrics["train/accuracy"] /= combined_metrics["train/epoch_size"] |
|
combined_metrics["train/loss"] /= combined_metrics["train/epoch_size"] |
|
combined_metrics["validation/accuracy"] /= combined_metrics["validation/epoch_size"] |
|
combined_metrics["validation/loss"] /= combined_metrics["validation/epoch_size"] |
|
|
|
if rank == 0: |
|
for key, value in combined_metrics.items(): |
|
writer.add_scalar(key, value, global_step=epoch) |
|
|
|
if combined_metrics["validation/accuracy"] > best_validation_accuracy: |
|
best_validation_accuracy = combined_metrics["validation/accuracy"] |
|
|
|
model_to_save = model.module if hasattr(model, 'module') else model |
|
torch.save(dict( |
|
epoch=epoch, |
|
model_state_dict=model_to_save.state_dict(), |
|
optimizer_state_dict=optimizer.state_dict(), |
|
args=args |
|
), |
|
os.path.join(logdir, "best-model.pt") |
|
) |
|
|
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser() |
|
|
|
parser.add_argument('--max-epochs', type=int, default=None) |
|
parser.add_argument('--device', type=str, default=None) |
|
parser.add_argument('--batch-size', type=int, default=24) |
|
parser.add_argument('--max-sequence-length', type=int, default=128) |
|
parser.add_argument('--random-sequence-length', action='store_true') |
|
parser.add_argument('--epoch-size', type=int, default=None) |
|
parser.add_argument('--seed', type=int, default=None) |
|
parser.add_argument('--data-dir', type=str, default='data') |
|
parser.add_argument('--real-dataset', type=str, default='webtext') |
|
parser.add_argument('--fake-dataset', type=str, default='xl-1542M-k40') |
|
parser.add_argument('--token-dropout', type=float, default=None) |
|
|
|
parser.add_argument('--large', action='store_true', help='use the roberta-large model instead of roberta-base') |
|
parser.add_argument('--learning-rate', type=float, default=2e-5) |
|
parser.add_argument('--weight-decay', type=float, default=0) |
|
args = parser.parse_args() |
|
|
|
nproc = int(subprocess.check_output([sys.executable, '-c', "import torch;" |
|
"print(torch.cuda.device_count() if torch.cuda.is_available() else 1)"])) |
|
if nproc > 1: |
|
print(f'Launching {nproc} processes ...', file=sys.stderr) |
|
|
|
os.environ["MASTER_ADDR"] = '127.0.0.1' |
|
os.environ["MASTER_PORT"] = str(29500) |
|
os.environ['WORLD_SIZE'] = str(nproc) |
|
os.environ['OMP_NUM_THREAD'] = str(1) |
|
subprocesses = [] |
|
|
|
for i in range(nproc): |
|
os.environ['RANK'] = str(i) |
|
os.environ['LOCAL_RANK'] = str(i) |
|
process = Process(target=run, kwargs=vars(args)) |
|
process.start() |
|
subprocesses.append(process) |
|
|
|
for process in subprocesses: |
|
process.join() |
|
else: |
|
run(**vars(args)) |
|
|