| import torch |
| from omegaconf import OmegaConf |
| from swim.utils import instantiate_from_config |
| from torchinfo import summary |
| from swim.modules.dataset import SwimDataModule |
| from lightning import Trainer |
| from lightning.pytorch.loggers import WandbLogger |
|
|
| torch.set_float32_matmul_precision("medium") |
|
|
| config = OmegaConf.load("configs/autoencoder/autoencoder_kl_32x32x4.yaml") |
|
|
| model = instantiate_from_config(config.model) |
| model.learning_rate = config.model.base_learning_rate |
|
|
| datamodule = SwimDataModule( |
| root_dir="/cm/shared/ninhnq3/datasets/swim_data", batch_size=2, img_size=512 |
| ) |
|
|
| logger = WandbLogger(project="swim", name="autoencoder_kl") |
|
|
| trainer = Trainer(max_epochs=10, devices=[0], logger=logger, log_every_n_steps=10) |
| trainer.fit(model, datamodule) |
|
|
