diff --git a/check.py b/check.py index bd38e4e..29d2e8d 100644 --- a/check.py +++ b/check.py @@ -86,7 +86,11 @@ if __name__ == "__main__": name = f"out/v{v}" # ckpt = f"/teamspace/jobs/{name}/work/colors/lightning_logs/version_2/checkpoints/epoch=999-step=8000.ckpt" ckpt_path = f"/teamspace/studios/this_studio/colors/lightning_logs/version_{v}/checkpoints/*.ckpt" - ckpt = glob.glob(ckpt_path)[-1] - print(f"Generating image for checkpoint: {ckpt}") - create_circle(ckpt, fname=name) + ckpt = glob.glob(ckpt_path) + if len(ckpt) > 0: + ckpt = ckpt[-1] + print(f"Generating image for checkpoint: {ckpt}") + create_circle(ckpt, fname=name) + else: + print(f"No checkpoint found for version {v}") # make_image(ckpt, fname=name + "b", color=False) diff --git a/makefile b/makefile index 7da0b5d..6e25865 100644 --- a/makefile +++ b/makefile @@ -4,7 +4,7 @@ lint: flake8 --ignore E501,W503 . test: - python main.py --alpha 2 --lr 2e-4 --max_epochs 200 + python main.py --alpha 4 --lr 2e-4 --max_epochs 200 search: python search.py diff --git a/model.py b/model.py index cc358a8..8dd3b59 100644 --- a/model.py +++ b/model.py @@ -5,70 +5,106 @@ from torch.optim.lr_scheduler import ReduceLROnPlateau from losses import calculate_separation_loss, preservation_loss +# 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): 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), + # Neural network layers + self.network = nn.Sequential( + nn.Linear(5, 64), # Input layer + nn.Tanh(), + nn.Linear(64, 128), + nn.Tanh(), + nn.Linear(128, 128), + nn.Tanh(), + nn.Linear(128, 64), + nn.Tanh(), + nn.Linear(64, 1), # Output layer ) def forward(self, x): - x = self.layers(x) - x = (torch.sin(x) + 1) / 2 + # Pass the input through the network + x = self.network(x) + # Circular mapping + # x = (torch.sin(x) + 1) / 2 + x = torch.sigmoid(x) return x - # class ColorTransformerModel(pl.LightningModule): - # def __init__(self, alpha, learning_rate): - # super().__init__() - # self.save_hyperparameters() - - # # Embedding layer to expand the input dimensions - # self.embedding = nn.Linear(3, 128) - - # # Transformer block - # transformer_layer = nn.TransformerEncoderLayer( - # d_model=128, nhead=4, dim_feedforward=512, dropout=0.1 - # ) - # self.transformer_encoder = nn.TransformerEncoder( - # transformer_layer, num_layers=3 - # ) - - # # Final linear layer to map back to 1D space - # self.final_layer = nn.Linear(128, 1) - - # 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) - - # # 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) - - # return x - def training_step(self, batch, batch_idx): inputs, labels = batch # x are the RGB inputs, labels are the strings outputs = self.forward(inputs) @@ -84,11 +120,12 @@ class ColorTransformerModel(pl.LightningModule): return loss def configure_optimizers(self): - optimizer = torch.optim.AdamW( - self.parameters(), lr=self.hparams.learning_rate, + optimizer = torch.optim.SGD( + self.parameters(), + lr=self.hparams.learning_rate, ) lr_scheduler = ReduceLROnPlateau( - optimizer, mode="min", factor=0.1, patience=10, verbose=True + optimizer, mode="min", factor=0.05, patience=10, cooldown=20, verbose=True ) return { "optimizer": optimizer, diff --git a/utils.py b/utils.py index 0c41784..448dcc4 100644 --- a/utils.py +++ b/utils.py @@ -1,21 +1,22 @@ import torch -def preprocess_data(data): +def preprocess_data(data, skip=False): # Assuming 'data' is a tensor of shape [n_samples, 3] - - # Compute argmin and argmax for each row - argmin_values = torch.argmin(data, dim=1, keepdim=True).float() - argmax_values = torch.argmax(data, dim=1, keepdim=True).float() - - # Normalize or scale argmin and argmax if necessary - # For example, here I am just dividing by the number of features - argmin_values /= data.shape[1] - 1 - argmax_values /= data.shape[1] - 1 - - # Concatenate the argmin and argmax values to the original data - new_data = torch.cat((data, argmin_values, argmax_values), dim=1) - + if not skip: + # Compute argmin and argmax for each row + argmin_values = torch.argmin(data, dim=1, keepdim=True).float() + argmax_values = torch.argmax(data, dim=1, keepdim=True).float() + + # Normalize or scale argmin and argmax if necessary + # For example, here I am just dividing by the number of features + argmin_values /= data.shape[1] - 1 + argmax_values /= data.shape[1] - 1 + + # Concatenate the argmin and argmax values to the original data + new_data = torch.cat((data, argmin_values, argmax_values), dim=1) + else: + new_data = data return new_data