From 88c8cde9f6fd4a5f68d58794be4272f9cb21b55f Mon Sep 17 00:00:00 2001 From: "Michael Pilosov, PhD" Date: Sat, 27 Jan 2024 07:27:57 +0000 Subject: [PATCH] total overhaul of model --- dataloader.py | 6 ++- datamodule.py | 77 +++++++++++++++++++++++++++++ main.py | 5 +- model.py | 132 ++++++++++---------------------------------------- newmain.py | 28 +++++++++++ 5 files changed, 138 insertions(+), 110 deletions(-) create mode 100644 datamodule.py create mode 100644 newmain.py diff --git a/dataloader.py b/dataloader.py index 63120ac..75a779b 100644 --- a/dataloader.py +++ b/dataloader.py @@ -4,7 +4,7 @@ from torch.utils.data import DataLoader, TensorDataset from utils import extract_colors, preprocess_data -def create_dataloader(N: int = 1e8, skip: bool = True, **kwargs): +def create_random_dataloader(N: int = 1e8, skip: bool = True, **kwargs): rgb_tensor = torch.rand((int(N), 3), dtype=torch.float32) rgb_tensor = preprocess_data(rgb_tensor, skip=skip) # Creating a dataset and data loader @@ -36,7 +36,9 @@ def create_named_dataloader(N: int = 0, skip: bool = True, **kwargs): if __name__ == "__main__": batch_size = 4 - train_dataloader = create_dataloader(N=1e6, batch_size=batch_size, shuffle=True) + train_dataloader = create_random_dataloader( + N=1e6, batch_size=batch_size, shuffle=True + ) print(len(train_dataloader.dataset)) train_dataloader_with_names = create_named_dataloader( batch_size=batch_size, shuffle=True diff --git a/datamodule.py b/datamodule.py new file mode 100644 index 0000000..5d9b3ba --- /dev/null +++ b/datamodule.py @@ -0,0 +1,77 @@ +import lightning as L +import torch +from matplotlib.colors import rgb_to_hsv +from torch.utils.data import DataLoader + +from utils import extract_colors, preprocess_data + + +class ColorDataModule(L.LightningDataModule): + def __init__(self, val_size: int = 10_000, train_size=0, batch_size: int = 32): + super().__init__() + self.val_size = val_size + self.batch_size = batch_size + self.train_size = train_size + + def prepare_data(self): + # no state. called from main process. + pass + + def setup(self, stage: str): + rgb_tensor, xkcd_color_names = extract_colors() + rgb_tensor = preprocess_data(rgb_tensor, skip=True) + self.xkcd_colors = [ + (rgb_tensor[i], xkcd_color_names[i].replace("xkcd:", "")) + for i in range(len(rgb_tensor)) + ] + if self.train_size > 0: + train_rgb = torch.rand((int(self.val_size), 3), dtype=torch.float32) + train_rgb = preprocess_data(train_rgb, skip=True) + self.train_colors = [ + ( + train_rgb[i], + torch.tensor(rgb_to_hsv(train_rgb[i])[:, 0], dtype=torch.float32), + ) + for i in range(len(train_rgb)) + ] + + val_rgb = torch.rand((int(self.val_size), 3), dtype=torch.float32) + val_rgb = preprocess_data(val_rgb, skip=True) + self.random_colors = [ + ( + val_rgb[i], + torch.tensor(rgb_to_hsv(val_rgb[i])[:, 0], dtype=torch.float32), + ) + for i in range(len(val_rgb)) + ] + + # Assign train/val datasets for use in dataloaders + if stage == "fit": + self.color_val = self.random_colors + if self.train_size > 0: + self.color_train = self.train_colors + else: + self.color_train = self.xkcd_colors + + # Assign test dataset for use in dataloader(s) + if stage == "test": + self.color_test = self.random_colors + + if stage == "predict": # for visualizing + self.color_predict = self.xkcd_colors + + def train_dataloader(self): + return DataLoader(self.color_train, batch_size=self.batch_size) + + def val_dataloader(self): + return DataLoader(self.color_val, batch_size=self.batch_size) + + def test_dataloader(self): + return DataLoader(self.color_test, batch_size=self.batch_size) + + def predict_dataloader(self): + return DataLoader(self.color_predict, batch_size=self.batch_size) + + def teardown(self, stage: str): + # Used to clean-up when the run is finished + pass diff --git a/main.py b/main.py index 2fe83f9..aa71a40 100644 --- a/main.py +++ b/main.py @@ -79,7 +79,7 @@ if __name__ == "__main__": # named_data_loader also has grayscale extras. TODO: remove unnamed train_dataloader = create_dataloader( # N=1e5, - skip=False, + skip=True, batch_size=args.bs, shuffle=True, num_workers=args.num_workers, @@ -90,6 +90,9 @@ if __name__ == "__main__": learning_rate=args.lr, batch_size=args.bs, width=args.width, + bias=False, + transform="relu", + depth=1, ) # Initialize model with parsed arguments diff --git a/model.py b/model.py index e24415f..67fd130 100644 --- a/model.py +++ b/model.py @@ -1,4 +1,4 @@ -import pytorch_lightning as pl +import lightning as L import torch import torch.nn as nn from torch.optim.lr_scheduler import ReduceLROnPlateau @@ -6,118 +6,36 @@ from torch.optim.lr_scheduler import ReduceLROnPlateau from losses import calculate_separation_loss, preservation_loss # noqa: F401 from utils import PURE_HSV, PURE_RGB -# class ColorTransformerModel(pl.LightningModule): -# def __init__(self, params): -# super().__init__() -# self.save_hyperparameters(params) -# # Model layers -# self.layers = nn.Sequential( -# nn.Linear(5, 128, bias=False), -# nn.Linear(128, 3, bias=False), -# nn.ReLU(), -# nn.Linear(3, 64, bias=False), -# nn.Linear(64, 128, bias=False), -# nn.Linear(128, 256, bias=False), -# nn.Linear(256, 128, bias=False), -# nn.ReLU(), -# nn.Linear(128, 1, bias=False), -# ) - -# def forward(self, x): -# x = self.layers(x) -# x = (torch.sin(x) + 1) / 2 -# return x - -# class ColorTransformerModel(pl.LightningModule): -# def __init__(self, params): -# super().__init__() -# self.save_hyperparameters(params) - -# # Embedding layer to expand the input dimensions -# self.embedding = nn.Linear(3, 128, bias=False) - -# # Transformer encoder-decoder -# encoder = nn.TransformerEncoderLayer( -# d_model=128, nhead=4, dim_feedforward=512, dropout=0.3 -# ) -# self.transformer_encoder = nn.TransformerEncoder( -# encoder, num_layers=3 -# ) -# # lower dimensionality decoder -# decoder = nn.TransformerDecoderLayer( -# d_model=128, nhead=4, dim_feedforward=512, dropout=0.3 -# ) -# self.transformer_decoder = nn.TransformerDecoder( -# decoder, num_layers=3 -# ) - -# # Final linear layer to map back to 1D space -# self.final_layer = nn.Linear(128, 1, bias=False) - -# def forward(self, x): -# # Embedding the input -# x = self.embedding(x) - -# # Adjusting the shape for the transformer -# x = x.unsqueeze(1) # Adding a fake sequence dimension - -# # Passing through the transformer -# x = self.transformer_encoder(x) - -# # Passing through the decoder -# x = self.transformer_decoder(x, memory=x) - -# # Reshape back to original shape -# x = x.squeeze(1) - -# # Final linear layer -# x = self.final_layer(x) - -# # Apply sigmoid activation to ensure output is in (0, 1) -# # x = torch.sigmoid(x) -# x = (torch.sin(x) + 1) / 2 -# return x - - -class ColorTransformerModel(pl.LightningModule): - def __init__(self, params): +class ColorTransformerModel(L.LightningModule): + def __init__( + self, + transform: str = "relu", + width: int = 128, + depth: int = 1, + bias: bool = False, + ): super().__init__() - self.save_hyperparameters(params) - # self.a = nn.Sequential( - # nn.Linear(3, 3, bias=False), - # nn.ReLU(), - # nn.Linear(3, 3, bias=False), - # nn.ReLU(), - # nn.Linear(3, 1, bias=False), - # nn.ReLU(), - # ) - # self.b = nn.Sequential( - # nn.Linear(3, 3, bias=False), - # nn.ReLU(), - # nn.Linear(3, 3, bias=False), - # nn.ReLU(), - # nn.Linear(3, 1, bias=False), - # nn.ReLU(), - # ) - # Neural network layers + self.save_hyperparameters() + if self.hparams.transform.lower() == "tanh": + t = nn.Tanh + elif self.hparams.transform.lower() == "relu": + t = nn.ReLU + + w = self.hparams.width + d = self.hparams.depth + bias = self.hparams.bias + midlayers = [nn.Linear(w, w, bias=bias), t()] * d self.network = nn.Sequential( - nn.Linear(5, 64), - nn.Tanh(), - nn.Linear(64, self.hparams.width), - nn.Tanh(), - nn.Linear(self.hparams.width, 3), - nn.Tanh(), - nn.Linear(3, 1), + nn.Linear(3, w, bias=bias), + t(), + *midlayers, + nn.Linear(w, 3, bias=bias), + t(), + nn.Linear(3, 1, bias=bias), ) def forward(self, x): - # Pass the input through the network - # a = self.a(x) - # b = self.b(x) - # a = torch.sigmoid(a) - # b = torch.sigmoid(b) - # x = torch.cat([x, a, b], dim=-1) x = self.network(x) # Circular mapping # x = (torch.sin(x) + 1) / 2 diff --git a/newmain.py b/newmain.py new file mode 100644 index 0000000..eec05a7 --- /dev/null +++ b/newmain.py @@ -0,0 +1,28 @@ +from lightning.pytorch.cli import LightningCLI + +# from callbacks import SaveImageCallback +from datamodule import ColorDataModule +from model import ColorTransformerModel + + +def cli_main(): + cli = LightningCLI(ColorTransformerModel, ColorDataModule) # noqa: F841 + # note: don't call fit!! + + +if __name__ == "__main__": + cli_main() + # note: it is good practice to implement the CLI in a function and call it in the main if block + + +# save_img_callback = SaveImageCallback( +# save_interval=0, +# final_dir="out", +# ) + +# trainer = pl.Trainer( +# callbacks=[save_img_callback], +# ) + +# data = ColorDataModule() +# trainer.fit(model, data)