in detector/train.py [0:0]
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")
)