File size: 1,303 Bytes
08c9919 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import AGen_model
from lib.datasets.videos_dataset import create_video_dataloader
import hydra
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
import os
import glob
@hydra.main(config_path="confs", config_name="base")
def main(opt):
pl.seed_everything(42)
print("Working dir:", os.getcwd())
checkpoint_callback = pl.callbacks.ModelCheckpoint(
dirpath="checkpoints/",
filename="{epoch:04d}-{loss}",
save_on_train_epoch_end=True,
save_last=True)
logger = WandbLogger(project=opt.project_name, name=f"{opt.exp}")
AGen_trainer = pl.Trainer(
gpus=2,
accelerator="gpu",
callbacks=[checkpoint_callback],
max_epochs=8000,
check_val_every_n_epoch=50,
logger=logger,
log_every_n_steps=1,
num_sanity_val_steps=0
)
model = AGen_model(opt)
trainset = create_video_dataloader(opt.videos_dataset.train)
validset = create_video_dataloader(opt.videos_dataset.valid)
if opt.model.is_continue == True:
checkpoint = sorted(glob.glob("checkpoints/*.ckpt"))[-1]
AGen_trainer.fit(model, trainset, validset, ckpt_path=checkpoint)
else:
AGen_trainer.fit(model, trainset, validset)
if __name__ == '__main__':
main() |