Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132from argparse import ArgumentParserfrom datetime import datetimeimport osfrom typing import Optional
import lightning as pylfrom lightning.pytorch.loggers import WandbLoggerfrom lightning.pytorch.callbacks import ModelCheckpointimport torchfrom torchdata.stateful_dataloader import StatefulDataLoader
from scripts.training.data import FollowDatasetfrom scripts.training.models import FollowEmbedModulefrom scripts.utils import get_logger
logger = get_logger(__name__)
PROJECT = "goodposts-followers"SAVE_DIR = "./data/training"NAME = "following-embedding"
def main(run_id: Optional[str], model_checkpoint: Optional[str], dataset_path: str): if run_id is None: run_id = NAME + str(int(datetime.now().timestamp()))
logger.info("Loading training dataset") train_dataset = FollowDataset( dataset_path=dataset_path, split="training", negative_sample_chance=0.5, )
logger.info("Loading validation dataset") val_dataset = FollowDataset( dataset_path=dataset_path, split="validation", negative_sample_chance=0.0, )
train_dataloader = StatefulDataLoader( train_dataset, pin_memory=True, collate_fn=train_dataset.collate_rows, batch_size=1024, num_workers=7, shuffle=True, )
val_dataloader = StatefulDataLoader( val_dataset, pin_memory=True, collate_fn=val_dataset.collate_rows, batch_size=1024, num_workers=7, )
logger.info("Instantiating embedding model") wandb_logger = WandbLogger( project=PROJECT, save_dir=SAVE_DIR, name=NAME, id=run_id, resume="allow", )
checkpoint_callback = ModelCheckpoint( dirpath=os.path.join(SAVE_DIR, "checkpoints", run_id), filename="{epoch}-{step}", save_last=True, save_top_k=1, monitor="val_loss", )
if model_checkpoint: embedding = FollowEmbedModule.load_from_checkpoint( checkpoint_path=model_checkpoint ) else: embedding = FollowEmbedModule( num_embeds=train_dataset.num_users(), embed_dim=256, learning_rate=1e-1, )
torch.set_float32_matmul_precision("medium") trainer = pyl.Trainer( logger=wandb_logger, callbacks=[checkpoint_callback], accelerator="gpu", devices=1, precision="bf16", gradient_clip_val=0.5, max_epochs=10, val_check_interval=0.5, ) trainer.fit( model=embedding, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader, ckpt_path=model_checkpoint, )
if __name__ == "__main__": parser = ArgumentParser() parser.add_argument( "-r", "--run-id", dest="run_id", help="ID for training run", ) parser.add_argument( "-m", "--model-checkpoint", dest="model_checkpoint", help="Path to checkpoint with run data", ) parser.add_argument( "-d", "--dataset-path", dest="dataset_path", help="Path to dataset directory", )
args = parser.parse_args() main( run_id=args.run_id, model_checkpoint=args.model_checkpoint, dataset_path=args.dataset_path, )