39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
import torch
|
|
from lightning import Trainer
|
|
from ThermalSolver.module import GraphSolver
|
|
from ThermalSolver.data_module import GraphDataModule
|
|
from ThermalSolver.model.local_gno import GatingGNO
|
|
|
|
|
|
def main():
|
|
trainer = Trainer(
|
|
max_epochs=50, accelerator="cuda", devices=1, accumulate_grad_batches=3
|
|
)
|
|
data_module = GraphDataModule(
|
|
hf_repo="SISSAmathLab/thermal-conduction",
|
|
split_name="2000",
|
|
train_size=0.8,
|
|
val_size=0.1,
|
|
test_size=0.1,
|
|
batch_size=10,
|
|
)
|
|
data_module.prepare_data()
|
|
data_module.setup("fit")
|
|
model = GatingGNO(
|
|
x_ch_node=1, f_ch_node=1, hidden=16, layers=2, edge_ch=3, out_ch=1
|
|
)
|
|
solver = GraphSolver(model, unrolling_steps=64)
|
|
|
|
trainer.fit(
|
|
solver,
|
|
train_dataloaders=data_module.train_dataloader(),
|
|
val_dataloaders=data_module.val_dataloader(),
|
|
)
|
|
data_module.setup("test")
|
|
trainer.test(solver, dataloaders=data_module.test_dataloader())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
torch.set_float32_matmul_precision("medium")
|
|
main()
|