Skip to content

Commit

Permalink
comply with ruff linter
Browse files Browse the repository at this point in the history
  • Loading branch information
@kmario23
kmario23 committed Oct 20, 2024
1 parent be93f68 commit dc19ad9
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 31 deletions.
21 changes: 9 additions & 12 deletions pdebench/models/fno/fno.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@

class SpectralConv1d(nn.Module):
def __init__(self, in_channels, out_channels, modes1):
super(SpectralConv1d, self).__init__()
super().__init__()

"""
1D Fourier layer. It does FFT, linear transform, and Inverse FFT.
Expand Down Expand Up @@ -77,13 +77,12 @@ def forward(self, x):
)

# Return to physical space
x = torch.fft.irfft(out_ft, n=x.size(-1))
return x
return torch.fft.irfft(out_ft, n=x.size(-1))


class FNO1d(nn.Module):
def __init__(self, num_channels, modes=16, width=64, initial_step=10):
super(FNO1d, self).__init__()
super().__init__()

"""
The overall network. It contains 4 layers of the Fourier layer.
Expand Down Expand Up @@ -155,7 +154,7 @@ def forward(self, x, grid):

class SpectralConv2d_fast(nn.Module):
def __init__(self, in_channels, out_channels, modes1, modes2):
super(SpectralConv2d_fast, self).__init__()
super().__init__()

"""
2D Fourier layer. It does FFT, linear transform, and Inverse FFT.
Expand Down Expand Up @@ -210,13 +209,12 @@ def forward(self, x):
)

# Return to physical space
x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
return x
return torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))


class FNO2d(nn.Module):
def __init__(self, num_channels, modes1=12, modes2=12, width=20, initial_step=10):
super(FNO2d, self).__init__()
super().__init__()

"""
The overall network. It contains 4 layers of the Fourier layer.
Expand Down Expand Up @@ -297,7 +295,7 @@ def forward(self, x, grid):

class SpectralConv3d(nn.Module):
def __init__(self, in_channels, out_channels, modes1, modes2, modes3):
super(SpectralConv3d, self).__init__()
super().__init__()

"""
3D Fourier layer. It does FFT, linear transform, and Inverse FFT.
Expand Down Expand Up @@ -392,15 +390,14 @@ def forward(self, x):
)

# Return to physical space
x = torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1)))
return x
return torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1)))


class FNO3d(nn.Module):
def __init__(
self, num_channels, modes1=8, modes2=8, modes3=8, width=20, initial_step=10
):
super(FNO3d, self).__init__()
super().__init__()

"""
The overall network. It contains 4 layers of the Fourier layer.
Expand Down
38 changes: 19 additions & 19 deletions pdebench/models/fno/train.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,20 +2,19 @@

import pickle
from timeit import default_timer
from pathlib import Path

import numpy as np
import torch
from torch import nn

# torch.manual_seed(0)
# np.random.seed(0)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from pdebench.models.fno.fno import FNO1d, FNO2d, FNO3d
from pdebench.models.fno.utils import FNODatasetMult, FNODatasetSingle
from pdebench.models.metrics import metrics

# torch.manual_seed(0)
# np.random.seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def run_training(
if_training,
Expand Down Expand Up @@ -48,9 +47,9 @@ def run_training(
base_path="../data/",
training_type="autoregressive",
):
print(
f"Epochs = {epochs}, learning rate = {learning_rate}, scheduler step = {scheduler_step}, scheduler gamma = {scheduler_gamma}"
)
# print(
# f"Epochs = {epochs}, learning rate = {learning_rate}, scheduler step = {scheduler_step}, scheduler gamma = {scheduler_gamma}"
# )

################################################################
# load data
Expand All @@ -59,7 +58,7 @@ def run_training(
if single_file:
# filename
model_name = flnm[:-5] + "_FNO"
print("FNODatasetSingle")
# print("FNODatasetSingle")

# Initialize the dataset and dataloader
train_data = FNODatasetSingle(
Expand All @@ -84,7 +83,7 @@ def run_training(
# filename
model_name = flnm + "_FNO"

print("FNODatasetMult")
# print("FNODatasetMult")
train_data = FNODatasetMult(
flnm,
reduced_resolution=reduced_resolution,
Expand Down Expand Up @@ -114,7 +113,7 @@ def run_training(

_, _data, _ = next(iter(val_loader))
dimensions = len(_data.shape)
print("Spatial Dimension", dimensions - 3)
# print("Spatial Dimension", dimensions - 3)
if dimensions == 4:
model = FNO1d(
num_channels=num_channels,
Expand Down Expand Up @@ -147,7 +146,7 @@ def run_training(
model_path = model_name + ".pt"

total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total parameters = {total_params}")
# print(f"Total parameters = {total_params}")

optimizer = torch.optim.Adam(
model.parameters(), lr=learning_rate, weight_decay=1e-4
Expand Down Expand Up @@ -184,14 +183,15 @@ def run_training(
t_max,
initial_step=initial_step,
)
pickle.dump(errs, open(model_name + ".pickle", "wb"))
with Path(model_name + ".pickle").open("wb") as pb:
pickle.dump(errs, pb)

return

# If desired, restore the network by loading the weights saved in the .pt
# file
if continue_training:
print("Restoring model (that is the network's weights) from file...")
# print("Restoring model (that is the network's weights) from file...")
checkpoint = torch.load(model_path, map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
model.to(device)
Expand Down Expand Up @@ -339,11 +339,11 @@ def run_training(

t2 = default_timer()
scheduler.step()
print(
"epoch: {0}, loss: {1:.5f}, t2-t1: {2:.5f}, trainL2: {3:.5f}, testL2: {4:.5f}".format(
ep, loss.item(), t2 - t1, train_l2_full, val_l2_full
)
)
# print(
# "epoch: {0}, loss: {1:.5f}, t2-t1: {2:.5f}, trainL2: {3:.5f}, testL2: {4:.5f}".format(
# ep, loss.item(), t2 - t1, train_l2_full, val_l2_full
# )
# )


if __name__ == "__main__":
Expand Down

0 comments on commit dc19ad9

Please sign in to comment.