diff --git a/README.md b/README.md
index e0b09f8..281d9af 100644
--- a/README.md
+++ b/README.md
@@ -19,44 +19,12 @@ Then, you can import the module `my_pytorch_kit`.
## Usage
This package revolves around the [`BaseModel`](my_pytorch_kit/model/models.py), [`Trainer`](my_pytorch_kit/train/train.py) and [`Evaluator`](my_pytorch_kit/evaluation/evaluation.py) classes, which are extended to model, train and evaluate a model respectively.
-```mermaid
-graph TD
- subgraph "Usage Workflow"
- A["Define Model
(extends BaseModel)"]
- B["Define Data
(Dataset / DataLoader)"]
- G["Extras
(Optimizer, Tensorboard)"]
- C["Initialize Trainer
(extends Trainer)"]
- D["Initialize Evaluator
(extends Evaluator)"]
- E["trainer.train()"]
- F["evaluator.evaluate()"]
- G["Intitialize Tuner
(Hyperparameter Tuning)"]
- H["tuner.tune()"]
- I["Provided Architectures
(Classifier, AE, VAE, ...)"]
- end
- %% Define node relationships
- A --> C
- B --> C
- I --> A
- A --> D
- B --> D
- C --> E
- D --> F
- C --> G
- G --> H
+
+ 
+
- %% Style the nodes
- style A fill:#fbe,stroke:#333,stroke-width:2px
- style B fill:#fbe,stroke:#333,stroke-width:2px
- style G fill:#ffc,stroke:#333,stroke-width:2px
- style C fill:#cde,stroke:#333,stroke-width:2px
- style D fill:#cde,stroke:#333,stroke-width:2px
- style E fill:#cfc,stroke:#333,stroke-width:2px
- style F fill:#cfc,stroke:#333,stroke-width:2px
- style G fill:#cde,stroke:#333,stroke-width:2px
- style H fill:#cfc,stroke:#333,stroke-width:2px
-```
Furthermore, this package provides architecture implementations and modelling utilities.
Currently implemented architectures include:
diff --git a/docs/assets/vae_4_generation.png b/docs/assets/vae_4_generation.png
index 93cd45e..3547349 100644
Binary files a/docs/assets/vae_4_generation.png and b/docs/assets/vae_4_generation.png differ
diff --git a/examples/mnist/classifier/classifier.py b/examples/mnist/classifier/classifier.py
index 3062ece..ff4f384 100644
--- a/examples/mnist/classifier/classifier.py
+++ b/examples/mnist/classifier/classifier.py
@@ -58,7 +58,7 @@ def evaluate_batch(self, model, batch):
trainer.train(optimizer, **hparams)
- model.save_model("models/classifier.pt")
+ model.save_model("models/classifier.pt", hparams)
else:
model.load_model("models/classifier.pt")
diff --git a/examples/mnist/classifier/example.py b/examples/mnist/classifier/example.py
index adbe4d1..53a01cb 100644
--- a/examples/mnist/classifier/example.py
+++ b/examples/mnist/classifier/example.py
@@ -92,7 +92,7 @@ def accumulate_result(self, result):
def get_result(self):
return self.result
- def on_eval(self):
+ def on_eval(self, model):
self.batch_count = 0
#++++++++++++++ Here comes the main code ++++++++++++++++#
diff --git a/examples/mnist/classifier/tune.py b/examples/mnist/classifier/tune.py
index 75224fa..b840a55 100644
--- a/examples/mnist/classifier/tune.py
+++ b/examples/mnist/classifier/tune.py
@@ -85,7 +85,8 @@
num_search=num_search,
ranks_considered=ranks_considered,
check_multiplicant=check_multiplier,
- mode='dynamic')
+ mode='dynamic',
+ model_path='models/mnist_random_dynamic_')
result = evaluator.evaluate(model, test_loader)
print(f"Model accuracy on test set doing random dynamic search: {result*100:.2f}%")
diff --git a/examples/mnist/vae/vae.py b/examples/mnist/vae/vae.py
index c5c7a20..73662d0 100644
--- a/examples/mnist/vae/vae.py
+++ b/examples/mnist/vae/vae.py
@@ -20,25 +20,26 @@ def kl_annealing(epoch, epochs, model, **kwargs):
train = input("Train new model? (y/[n]) ").lower() == "y"
hparams = {
- "learning_rate": 1e-3,
+ "learning_rate": 3e-3,
"batch_size": 64,
"epochs": 15,
"optimizer_method": "Adam",
# "optimizer_kwargs": {"weight_decay": 1e-4},
"loss_func": nn.BCELoss(reduction="sum"),
"alpha": 1e-1,
- "beta": 1.5e-3,
- "classifier_loss_weight": 1.2,
+ "beta": 1.5e-1,
+ "classifier_loss_weight": 12,
"start_beta": 1e-3,
"final_beta": 4e-2,
"feature_space": (64, 7, 7),
- "latent_dim": 4,
+ "latent_dim": 2,
"sample_input_shape": (1, 1, 28, 28),
"classifier_num_layers": 3,
# "epoch_function": kl_annealing
}
model = ImageVAESemiSupervised(**hparams)
+ model.proper_weight_init()
# model = ImageVAE(**hparams)
torchsummary.summary(model, (1, 28, 28))
@@ -64,7 +65,10 @@ def kl_annealing(epoch, epochs, model, **kwargs):
model.save_model("models/vae.pt", hparams)
else:
- model.load_model("models/vae.pt")
+ path = input("Enter path to model (default: models/vae.pt) :")
+ if path == "":
+ path = "models/vae.pt"
+ model.load_model(path)
model.eval()
diff --git a/examples/mnist/vae/vae_semi_generate.py b/examples/mnist/vae/vae_semi_generate.py
index 5c3d897..40199bd 100644
--- a/examples/mnist/vae/vae_semi_generate.py
+++ b/examples/mnist/vae/vae_semi_generate.py
@@ -4,13 +4,25 @@
import os
from my_pytorch_kit.model.vae import ImageVAESemiSupervised
+from my_pytorch_kit.model.classifier import ImageClassifier
+from my_pytorch_kit.model.ae import ImageAE
-def generate_samples(model, latent_dim, num_samples=20):
+def generate_samples(model, latent_dim, num_samples=20, check_classifier=False, check_ae=False):
"""
Generates num_samples samples from the model,
plotting them in a grid with their labels
"""
+ classifier = None
+ if check_classifier:
+ classifier = ImageClassifier()
+ classifier.load_model("models/classifier.pt")
+
+ ae = None
+ if check_ae:
+ ae = ImageAE()
+ ae.load_model("models/ae.pt")
+
num_samples = int(num_samples)
n_rows = num_samples // 10
n_cols = 10
@@ -20,13 +32,40 @@ def generate_samples(model, latent_dim, num_samples=20):
z = torch.randn((num_samples, latent_dim))
images, labels = model.generate(z)
+ cl_labels = None
+ if classifier:
+ cl_labels = classifier(images)
+ cl_labels = torch.argmax(cl_labels, dim=1)
+
+ batch_recon_losses = None
+ if ae:
+ reconstruction = ae(images)
+ recon_losses = torch.nn.BCELoss(reduction="none")(reconstruction, images)
+ batch_recon_losses = recon_losses.mean(dim=(1, 2, 3))
+
plt.figure(figsize=(10, 10))
for i in range(num_samples):
plt.subplot(n_rows, n_cols, i + 1)
plt.imshow(images[i].view(28, 28).cpu().numpy(), cmap="Greys_r")
- label = torch.argmax(labels[i]).item()
+ label = f"{torch.argmax(labels[i]).item()}"
+ if cl_labels is not None:
+ label += f" ({cl_labels[i].item()})"
+ if recon_losses is not None:
+ label += f" ({batch_recon_losses[i].item():.2f})"
+
plt.title(label)
plt.axis("off")
+
+ extra_text = "Generator label"
+ if cl_labels is not None:
+ extra_text += " (Classifier label)"
+ if recon_losses is not None:
+ extra_text += " (Reconstruction loss)"
+ plt.text(-100, 50, extra_text, ha="center", va="center")
+
+ print(labels.shape, cl_labels.shape)
+ print(torch.nn.CrossEntropyLoss()(labels, cl_labels).item())
+
plt.show()
@@ -76,7 +115,7 @@ def generate_dataset(model, num_samples, latent_dim, std_divider=1.5, batch_size
model.load_model("models/vae_semi_3.pt")
- generate_samples(model, hparams["latent_dim"])
+ generate_samples(model, hparams["latent_dim"], check_classifier=True, check_ae=True)
gen_dataset = input("Generate dataset? (y/[n]): ").lower() == "y"
diff --git a/examples/mnist/vae/vae_semi_tune.py b/examples/mnist/vae/vae_semi_tune.py
new file mode 100644
index 0000000..87cdbdd
--- /dev/null
+++ b/examples/mnist/vae/vae_semi_tune.py
@@ -0,0 +1,148 @@
+
+import torch
+import torch.nn as nn
+import tqdm
+
+from my_pytorch_kit.model.vae import ImageVAESemiSupervised
+from my_pytorch_kit.train.tune import Tuner
+from my_pytorch_kit.evaluation.evaluation import Evaluator
+from my_pytorch_kit.model.classifier import ImageClassifier
+from my_pytorch_kit.model.ae import ImageAE
+from mnist.utils.mnist_utils import get_mnist_loaders
+from mnist.autoencoder.ae_plots import plot_latent, plot_reconstructions, plot_label_clusters
+
+
+class GenerationEvaluator(Evaluator):
+
+
+ def __init__(self,
+ latent_dim = 2,
+ num_samples = 1000,
+ batch_size = 64,
+ ce_weight = 1,
+ bce_weight = 2.5):
+ super().__init__()
+ self.batch_count = 0
+ self.result = 0
+
+ self.classifier = ImageClassifier()
+ self.classifier.load_model("models/classifier.pt")
+
+ self.ae = ImageAE()
+ self.ae.load_model("models/ae.pt")
+
+ self.latent_dim = latent_dim
+ self.num_samples = num_samples
+ self.batch_size = batch_size
+
+ self.bce = nn.BCELoss(reduction="mean")
+ self.ce = nn.CrossEntropyLoss(reduction="mean", label_smoothing=0.1)
+
+ self.ce_weight = ce_weight
+ self.bce_weight = bce_weight
+
+ def on_eval(self, model):
+
+ print("\nEvaluating...")
+ generated_count = 0
+
+ acc = 0
+
+ model.eval()
+ pbar = tqdm.tqdm(total=self.num_samples // self.batch_size + 1)
+ with torch.no_grad():
+ while generated_count < self.num_samples:
+
+ next_batch_size = (self.num_samples - generated_count) % self.batch_size
+ if next_batch_size == 0:
+ next_batch_size = self.batch_size
+
+
+ z = torch.randn((next_batch_size, self.latent_dim))
+ images, labels = model.generate(z)
+
+ cl_labels = self.classifier(images)
+ cl_labels = torch.argmax(cl_labels, dim=1)
+
+ ce_loss = self.ce(labels, cl_labels)
+
+
+ reconstruction = self.ae(images)
+ recon_losses = self.bce(reconstruction, images)
+
+ loss = self.bce_weight * recon_losses.item() + self.ce_weight * ce_loss.item()
+
+ acc += loss
+
+ generated_count += next_batch_size
+
+ pbar.update(1)
+ self.result = acc
+
+ def evaluate_batch(self, model, batch):
+ return 0
+
+ def accumulate_result(self, result):
+ return
+
+ def get_result(self):
+ return self.result
+
+def kl_annealing(epoch, epochs, model, **kwargs):
+ model.beta = model.beta + (kwargs["final_beta"] - kwargs["start_beta"]) / (epochs - kwargs["beta_start_delay"])
+
+
+if __name__ == '__main__':
+
+ hparams = {
+ "learning_rate": ((1e-4, 1e-2), "log"),
+ "batch_size": 64,
+ "epochs": 10,
+ "patience": 5,
+ "optimizer_method": "Adam",
+ # "optimizer_kwargs": {"weight_decay": 1e-4},
+ "loss_func": nn.BCELoss(reduction="sum"),
+ "alpha": ((0, 10), "float"),
+ "beta": ((0, 10), "float"),
+ "classifier_loss_weight": ((0, 10), "float"),
+ "start_beta": ((0, 1), "float"),
+ "final_beta": ((1, 10), "float"),
+ "beta_start_delay": ((4, 12), "int"),
+ "feature_space": ([(64, 7, 7)], "item"),
+ "latent_dim": 2,
+ "sample_input_shape": ([(1, 1, 28, 28)], "item"),
+ "classifier_num_layers": 3,
+ "use_scheduler": ((True, False, False), "item"),
+ "use_grad_clip": ((True, False, False), "item"),
+ "gamma": ((0.8, 0.95), "log"),
+ "epoch_function": ((kl_annealing, None, None), "item"),
+ }
+
+
+ train_loader, val_loader, test_loader = get_mnist_loaders(hparams["batch_size"])
+
+ evaluator = GenerationEvaluator()
+
+ tuner = Tuner(ImageVAESemiSupervised)
+
+ num_search = 500
+ ranks_considered = 15
+ check_multiplier = 3
+
+ # train model
+ model, best_config, results = tuner.tune(train_loader, val_loader, hparams,
+ evaluator=evaluator,
+ num_search=num_search,
+ ranks_considered=ranks_considered,
+ check_multiplicant=check_multiplier,
+ mode='dynamic', model_path='models/mnist_vae_semi_tuned_')
+
+ model.eval()
+
+ model.save_model("models/mnist_vae_semi_tuned_best.pt", best_config)
+
+ plot_reconstructions(model, test_loader)
+
+ distribution = plot_label_clusters(model.encoder, test_loader)
+
+ plot_latent(model.decoder, dist=distribution, stds=2)
diff --git a/examples/mnist/vae/vae_with_ae.py b/examples/mnist/vae/vae_with_ae.py
new file mode 100644
index 0000000..bc537e0
--- /dev/null
+++ b/examples/mnist/vae/vae_with_ae.py
@@ -0,0 +1,74 @@
+
+import torch.nn as nn
+import torchsummary
+
+from my_pytorch_kit.model.vae import ImageVAEwithAE
+from my_pytorch_kit.train.train import Trainer
+from my_pytorch_kit.train.optimizers import get_optimizer_total_optimizer
+from mnist.utils.mnist_utils import get_mnist_loaders
+from mnist.autoencoder.ae_plots import plot_latent, plot_reconstructions, plot_label_clusters
+
+
+def kl_annealing(epoch, epochs, model, **kwargs):
+ if epoch < kwargs["beta_start_delay"]:
+ return
+ model.beta = model.beta + (kwargs["final_beta"] - kwargs["start_beta"]) / (epochs - kwargs["beta_start_delay"])
+
+
+if __name__ == '__main__':
+
+ train = input("Train new model? (y/[n]) ").lower() == "y"
+
+ hparams = {
+ "learning_rate": 1e-3,
+ "batch_size": 64,
+ "epochs": 15,
+ "optimizer_method": "Adam",
+ # "optimizer_kwargs": {"weight_decay": 1e-4},
+ "loss_func": nn.BCELoss(reduction="sum"),
+ "alpha": 3,
+ "beta": 1,
+ "ae_reconstruction_loss_weight": 1,
+ "start_beta": 0,
+ "final_beta": 3e-1,
+ "beta_start_delay": 8,
+ "feature_space": (64, 7, 7),
+ "latent_dim": 2,
+ "sample_input_shape": (1, 1, 28, 28),
+ # "epoch_function": kl_annealing
+ }
+
+ model = ImageVAEwithAE(**hparams)
+
+ torchsummary.summary(model, (1, 28, 28))
+
+ train_loader, val_loader, test_loader = get_mnist_loaders(hparams["batch_size"])
+
+ optimizer = get_optimizer_total_optimizer(model, use_scheduler=False, use_grad_clip=False, **hparams)
+
+ trainer = Trainer(model, train_loader, val_loader, **hparams)
+
+
+ if train:
+ load_model = input("Load model? (y/[n]) ").lower() == "y"
+
+ if load_model:
+ path = input("Enter path to model (default: models/vae_with_ae.pt) :")
+ if path == "":
+ path = "models/vae_with_ae.pt"
+ model.load_model(path)
+
+ trainer.train(optimizer, **hparams)
+
+ model.save_model("models/vae_with_ae.pt", hparams)
+
+ else:
+ model.load_model("models/vae_with_ae.pt")
+
+ model.eval()
+
+ plot_reconstructions(model, test_loader)
+
+ distribution = plot_label_clusters(model.encoder, test_loader)
+
+ plot_latent(model.decoder, dist=distribution, stds=2)
diff --git a/examples/models/classifier.pt b/examples/models/classifier.pt
index ad7deb7..88b0caf 100644
Binary files a/examples/models/classifier.pt and b/examples/models/classifier.pt differ
diff --git a/examples/models/mnist_vae_semi_tuned.pt b/examples/models/mnist_vae_semi_tuned.pt
new file mode 100644
index 0000000..fad6628
Binary files /dev/null and b/examples/models/mnist_vae_semi_tuned.pt differ
diff --git a/examples/models/vae.pt b/examples/models/vae.pt
index 3c914af..8ca2963 100644
Binary files a/examples/models/vae.pt and b/examples/models/vae.pt differ
diff --git a/examples/models/vae_with_ae.pt b/examples/models/vae_with_ae.pt
new file mode 100644
index 0000000..75936cc
Binary files /dev/null and b/examples/models/vae_with_ae.pt differ
diff --git a/my_pytorch_kit/evaluation/evaluation.py b/my_pytorch_kit/evaluation/evaluation.py
index 1db62f0..3c080c4 100644
--- a/my_pytorch_kit/evaluation/evaluation.py
+++ b/my_pytorch_kit/evaluation/evaluation.py
@@ -57,9 +57,14 @@ def get_result(self) -> Any:
pass
- def on_eval(self):
+ def on_eval(self, model):
"""
Gets called before evaluation begins.
+
+ Parameters
+ ----------
+ model: torch.nn.Module
+ The model to evaluate.
"""
pass
@@ -80,7 +85,7 @@ def evaluate(self, model, data_loader) -> Any:
Any
The result metric.
"""
- self.on_eval()
+ self.on_eval(model)
model.eval()
for batch in data_loader:
result = self.evaluate_batch(model, batch)
diff --git a/my_pytorch_kit/evaluation/reconstruction.py b/my_pytorch_kit/evaluation/reconstruction.py
index 5b543d2..e08de46 100644
--- a/my_pytorch_kit/evaluation/reconstruction.py
+++ b/my_pytorch_kit/evaluation/reconstruction.py
@@ -36,7 +36,7 @@ def accumulate_result(self, result):
def get_result(self):
return self.acc
- def on_eval(self):
+ def on_eval(self, model):
if self.only_accumulate:
self.acc = []
else:
diff --git a/my_pytorch_kit/model/models.py b/my_pytorch_kit/model/models.py
index c3af817..ccdbc4c 100644
--- a/my_pytorch_kit/model/models.py
+++ b/my_pytorch_kit/model/models.py
@@ -3,7 +3,7 @@
from abc import abstractmethod
import os
import inspect
-from typing import Optional
+from typing import Optional, Union, Dict
class BaseModel(nn.Module):
"""
@@ -42,7 +42,7 @@ def __init_subclass__(cls, **kwargs):
)
@abstractmethod
- def calc_loss(self, batch, criterion, **kwargs) -> torch.Tensor:
+ def calc_loss(self, batch, criterion, **kwargs) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
"""
Calculates the loss for a given batch and criterion.
@@ -59,6 +59,11 @@ def calc_loss(self, batch, criterion, **kwargs) -> torch.Tensor:
-------
torch.Tensor
The loss for the batch.
+ OR
+ dict[str, torch.Tensor]
+ A dictionary of losses.
+ Losses should be keyed by name.
+ The total loss should be keyed by 'loss'.
"""
pass
diff --git a/my_pytorch_kit/model/vae.py b/my_pytorch_kit/model/vae.py
index 72f698c..5c1f36d 100644
--- a/my_pytorch_kit/model/vae.py
+++ b/my_pytorch_kit/model/vae.py
@@ -6,6 +6,7 @@
from my_pytorch_kit.model.models import BaseModel
from my_pytorch_kit.model.utils import ConvArchitect, AffineArchitect
+from my_pytorch_kit.model.ae import ImageAE
class ImageVAE(BaseModel):
@@ -46,6 +47,7 @@ def __init__(self,
self.alpha = alpha
self.beta = beta
+ self.loss_softmax = nn.Softmax()
self.encoder = VariationalEncoder(input_shape, encoder_num_layers, latent_dim, feature_space)
self.decoder = VariationalDecoder(input_shape, decoder_num_layers, latent_dim, feature_space)
@@ -67,9 +69,20 @@ def calc_loss(self, batch, criterion):
kl_div = self.kl_loss(self.params)
batch_size = x.size(0)
- loss = (self.alpha * reconstruction_loss / batch_size) + (self.beta * kl_div / batch_size)
+ alpha, beta = self.loss_softmax(torch.tensor([self.alpha, self.beta], dtype=float)).split(1)
- return loss
+ reconstruction_loss = reconstruction_loss * alpha / batch_size
+ kl_div = kl_div * beta / batch_size
+
+ loss = reconstruction_loss + kl_div
+
+ loss_dict = {
+ "loss": loss,
+ "reconstruction_loss": reconstruction_loss,
+ "kl_div": kl_div
+ }
+
+ return loss_dict
def kl_loss(self, params):
"""
@@ -101,6 +114,32 @@ def __init__(self,
classifier_activation = nn.PReLU(),
classifier_loss_weight = 1e-2,
**kwargs):
+ """
+ Initializes the Variational Autoencoder.
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the input image.
+ encoder_num_layers: int
+ Number of encoder layers.
+ decoder_num_layers: int
+ Number of decoder layers.
+ latent_dim: int
+ Dimension of the latent space.
+ feature_space: tuple
+ Shape of the output of the convolutional layers.
+ beta: float
+ Weight of the KL divergence loss.
+ classifier_num_layers: int
+ Number of layers in the classifier.
+ classifier_activation: nn.Module
+ Activation function for the classifier.
+ classifier_loss_weight: float
+ Weight of the classifier loss.
+ """
+
+
super().__init__(
input_shape = input_shape,
encoder_num_layers = encoder_num_layers,
@@ -117,6 +156,8 @@ def __init__(self,
self.classifier = AffineArchitect().build(latent_dim, num_classes, classifier_activation, classifier_num_layers)
self.cel = nn.CrossEntropyLoss()
+ self.proper_weight_init()
+
def forward(self, x):
params = self.encoder(x)
@@ -135,18 +176,134 @@ def generate(self, z):
def calc_loss(self, batch, criterion):
x, y = batch
x_hat, y_hat = self.forward(x)
+
reconstruction_loss = criterion(x_hat, x)
kl_div = self.kl_loss(self.params)
-
one_hot = nn.functional.one_hot(y, self.classifier_num_classes).float()
classifier_loss = self.cel(y_hat, one_hot)
batch_size = x.size(0)
- loss = (self.alpha * reconstruction_loss / batch_size) + (self.beta * kl_div / batch_size) + (self.classifier_loss_weight * classifier_loss / batch_size)
+ alpha, beta, classifier_loss_weight = self.loss_softmax(torch.tensor([self.alpha, self.beta, self.classifier_loss_weight], dtype=float)).split(1)
+
+ reconstruction_loss = (alpha * reconstruction_loss / batch_size)
+ kl_div = (beta * kl_div / batch_size)
+ classifier_loss = (classifier_loss_weight * classifier_loss / batch_size)
+
+ loss = reconstruction_loss + kl_div + classifier_loss
+
+ loss_dict = {
+ "loss": loss,
+ "reconstruction_loss": reconstruction_loss,
+ "kl_div": kl_div,
+ "classifier_loss": classifier_loss
+ }
+
+ return loss_dict
+
+
+class ImageVAEwithAE(ImageVAE):
+ """
+ Variational autoencoder, in which autoencoder reconstruction loss is added.
+ The autoencoder runs on the output of the VAE, ensuring that the VAE is forced
+ to produce images that can be represented in the latent space of the autoencoder.
+ """
+
+ def __init__(
+ self,
+ *,
+ input_shape = (1, 28, 28),
+ encoder_num_layers = 3,
+ decoder_num_layers = 3,
+ latent_dim = 2,
+ feature_space = (8, 3, 3),
+ alpha = 1,
+ beta = 1e-2,
+ ae_class = ImageAE,
+ ae_path = 'models/ae.pt',
+ ae_reconstruction_loss_weight = 1e-2,
+ ae_reconstruction_loss_criterion = nn.BCELoss(reduction="mean"),
+ **kwargs
+ ):
+ """
+ Initializes the Variational Autoencoder.
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the input image.
+ encoder_num_layers: int
+ Number of encoder layers.
+ decoder_num_layers: int
+ Number of decoder layers.
+ latent_dim: int
+ Dimension of the latent space.
+ feature_space: tuple
+ Shape of the output of the convolutional layers.
+ beta: float
+ Weight of the KL divergence loss.
+ ae_class: nn.Module
+ Autoencoder class
+ ae_path: str
+ Path to the autoencoder model
+ ae_reconstruction_loss_weight: float
+ Weight of the autoencoder reconstruction loss
+ ae_reconstruction_loss_criterion: nn.Module
+ Criterion for the autoencoder reconstruction loss
+ """
+ super().__init__(
+ input_shape = input_shape,
+ encoder_num_layers = encoder_num_layers,
+ decoder_num_layers = decoder_num_layers,
+ latent_dim = latent_dim,
+ feature_space = feature_space,
+ alpha = alpha,
+ beta = beta,
+ **kwargs
+ )
+
+ self.ae = ae_class()
+ self.ae.load_model(ae_path)
+ for param in self.ae.parameters():
+ param.requires_grad = False
+
+ self.ae_loss_weight = ae_reconstruction_loss_weight
+ self.ae_criterion = ae_reconstruction_loss_criterion
+
+ def forward(self, x):
+ params = self.encoder(x)
+ self.params = params
+ z = self.sampler(params)
+ x_hat = self.decoder(z)
+ x_hat = self.sigmoid(x_hat)
+ return x_hat
+
+ def calc_loss(self, batch, criterion):
+ x, _ = batch
+ x_hat = self.forward(x)
+
+ # get losses
+ reconstruction_loss = criterion(x_hat, x)
+ kl_div = self.kl_loss(self.params)
+ ae_recon_loss = self.ae_criterion(self.ae(x_hat), x_hat)
+
+ alpha, beta, ae_loss_weight = self.loss_softmax(torch.tensor([self.alpha, self.beta, self.ae_loss_weight], dtype=float)).split(1)
+ batch_size = x.size(0)
+
+ # then weight it
+ reconstruction_loss = reconstruction_loss * alpha / batch_size
+ kl_div = kl_div * beta / batch_size
+ ae_recon_loss = ae_recon_loss * ae_loss_weight / batch_size
- return loss
+ loss = reconstruction_loss + kl_div + ae_recon_loss
+ loss_dict = {
+ "loss": loss,
+ "reconstruction_loss": reconstruction_loss,
+ "kl_div": kl_div,
+ "ae_recon_loss": ae_recon_loss
+ }
+ return loss_dict
class VariationalEncoder(nn.Module):
"""
diff --git a/my_pytorch_kit/train/train.py b/my_pytorch_kit/train/train.py
index 4f53505..8be53db 100644
--- a/my_pytorch_kit/train/train.py
+++ b/my_pytorch_kit/train/train.py
@@ -56,9 +56,6 @@ def __init__(self,
if sample_input_shape is not None:
self.tb_logger.add_graph(self.model, torch.randn(*sample_input_shape))
- else:
- print("Sample input shape not given. Not adding graph to tensorboard logger.")
-
def train(self,
optimizer: TotalOptimizer,
@@ -132,7 +129,11 @@ def train(self,
training_loop = create_tqdm_bar(self.train_loader, desc=f'Training Epoch [{epoch + 1}/{epochs}]')
for train_iteration, batch in training_loop:
optimizer.zero_grad()
- loss = self.model.calc_loss(batch, loss_func)
+ returned_loss = self.model.calc_loss(batch, loss_func)
+ if type(returned_loss) is dict:
+ loss = returned_loss['loss']
+ else:
+ loss = returned_loss
loss.backward()
optimizer.step()
@@ -146,13 +147,24 @@ def train(self,
# Update the tensorboard logger.
self.add_tb_scalar(f'{name}/train_loss', loss.item(), epoch * len(self.train_loader) + train_iteration)
+ if type(returned_loss) is dict:
+ for loss_name, loss_value in returned_loss.items():
+ if loss_name != 'loss':
+ self.add_tb_scalar(f'{name}/loss_components/{loss_name}', loss_value.item() / loss.item(), epoch * len(self.train_loader) + train_iteration)
+
+
+
# VALIDATION
self.model.eval()
val_loop = create_tqdm_bar(self.val_loader, desc=f'Validation Epoch [{epoch + 1}/{epochs}]')
with torch.no_grad():
for val_iteration, batch in val_loop:
- loss = self.model.calc_loss(batch, loss_func)
+ returned_loss = self.model.calc_loss(batch, loss_func)
+ if type(returned_loss) is dict:
+ loss = returned_loss['loss']
+ else:
+ loss = returned_loss
validation_loss.append(loss.item())
last_cutoff_loss = np.mean(validation_loss[-loss_cutoff:])
@@ -162,6 +174,7 @@ def train(self,
# Update the tensorboard logger.
self.add_tb_scalar(f'{name}/val_loss', last_cutoff_loss, epoch * len(self.val_loader) + val_iteration)
+
# best val loss check
total_val_loss = np.mean(validation_loss)
if total_val_loss < best_val_loss:
diff --git a/my_pytorch_kit/train/tune.py b/my_pytorch_kit/train/tune.py
index c569627..f04d867 100644
--- a/my_pytorch_kit/train/tune.py
+++ b/my_pytorch_kit/train/tune.py
@@ -4,8 +4,11 @@
from typing import List, Dict, Tuple
import torch
from itertools import product
+import pprint
+import os
from my_pytorch_kit.train.train import Trainer
+from my_pytorch_kit.evaluation.evaluation import Evaluator
from my_pytorch_kit.train.optimizers import get_optimizer_total_optimizer
@@ -24,7 +27,13 @@ def __init__(self, model_class, trainer_class = Trainer):
self.trainer_class = trainer_class
- def tune(self, train_loader, val_loader, parameter_space, mode = 'grid', **kwargs):
+ def tune(self,
+ train_loader: torch.utils.data.DataLoader,
+ val_loader: torch.utils.data.DataLoader,
+ parameter_space: Dict[str, List] | Dict[str, Tuple[List, str]],
+ mode: str = 'grid',
+ evaluator: Evaluator = None,
+ **kwargs):
"""
Calls the appropriate search method and returns the best model.
@@ -39,19 +48,34 @@ def tune(self, train_loader, val_loader, parameter_space, mode = 'grid', **kwarg
Should be Dict[str, List] or Dict[str, Tuple[List, str]].
mode: str
'grid', 'random' or 'dynamic'
+ evaluator: Evaluator
+ Evaluator class. Will be used instead of validation loss if given.
+ Its evaluate() should return a single metric.
+
+ Returns
+ -------
+ best_model: torch.nn.Module
+ The model performing best on validation set
+ best_config: dict
+ The hyperparameter config that performed best
+ results: list
+ List of tuples (config, val_loss)
"""
if mode == 'grid':
- return self.grid_search(train_loader, val_loader, parameter_space)
+ return self.grid_search(train_loader, val_loader, parameter_space, evaluator=evaluator)
elif mode == 'random':
- return self.random_search(train_loader, val_loader, parameter_space, **kwargs)
+ return self.random_search(train_loader, val_loader, parameter_space, evaluator=evaluator, **kwargs)
elif mode == 'dynamic':
- return self.random_dynamic_search(train_loader, val_loader, parameter_space, **kwargs)
+ return self.random_dynamic_search(train_loader, val_loader, parameter_space, evaluator=evaluator, **kwargs)
else:
raise ValueError("Invalid mode: {}".format(mode))
- def grid_search(self, train_loader, val_loader, grid_search_spaces: Dict[str, List]):
+ def grid_search(self, train_loader,
+ val_loader,
+ grid_search_spaces: Dict[str, List],
+ evaluator: Evaluator = None):
"""
A simple grid search.
Searches all combinations of hyperparameters in grid_search_spaces.
@@ -66,6 +90,9 @@ def grid_search(self, train_loader, val_loader, grid_search_spaces: Dict[str, Li
Hyperparameter search spaces for grid search.
Specifies the possible values for each hyperparameter, e.g.
{'lr': [1e-3, 1e-4, 1e-5], 'batch_size': [32, 64, 128]}
+ evaluator: Evaluator
+ Evaluator class. Will be used instead of validation loss if given.
+ Its evaluate() should return a single metric.
Returns
-------
@@ -86,12 +113,15 @@ def grid_search(self, train_loader, val_loader, grid_search_spaces: Dict[str, Li
for instance in product(*grid_search_spaces.values()):
configs.append(dict(zip(grid_search_spaces.keys(), instance)))
- return self.find_best_config(configs, train_loader, val_loader)
+ return self.find_best_config(configs, train_loader, val_loader, evaluator)
- def random_search(self, train_loader, val_loader,
+ def random_search(self, train_loader,
+ val_loader,
random_search_spaces: Dict[str, Tuple[List, str]],
- num_search: int = 10, **kwargs):
+ num_search: int = 10,
+ evaluator: Evaluator = None,
+ **kwargs):
"""
Samples num_search hyper parameter sets within the provided search spaces
and returns the best model.
@@ -106,6 +136,9 @@ def random_search(self, train_loader, val_loader,
Hyperparameter search spaces for random search
num_search: int
Number of hyperparameter configs to sample (default: 10)
+ evaluator: Evaluator
+ Evaluator class. Will be used instead of validation loss if given.
+ Its evaluate() should return a single metric.
Returns
-------
@@ -126,12 +159,16 @@ def random_search(self, train_loader, val_loader,
for _ in range(num_search):
configs.append(self.random_search_spaces_to_config(random_search_spaces))
- return self.find_best_config(configs, train_loader, val_loader)
+ return self.find_best_config(configs, train_loader, val_loader, evaluator)
def random_dynamic_search(self, train_loader, val_loader,
random_search_spaces: Dict[str, Tuple[List, str]],
- num_search=10, ranks_considered=15,
- check_multiplicant=2, **kwargs):
+ num_search=10,
+ ranks_considered=15,
+ check_multiplicant=2,
+ evaluator: Evaluator = None,
+ model_path = None,
+ **kwargs):
"""
Samples num_search hyper parameter sets within the provided search space,
reducing the search space dynamically by looking at the best results,
@@ -152,6 +189,12 @@ def random_dynamic_search(self, train_loader, val_loader,
Number of best configs to consider
check_multiplicant: int
Multiplicant of ranks_considered
+ evaluator: Evaluator
+ Evaluator class. Will be used instead of validation loss if given.
+ Its evaluate() should return a single metric.
+ model_path : str
+ Saving path of the model. If given, saves all models in the ranking via
+ model_path.pt.
Returns
-------
@@ -163,19 +206,48 @@ def random_dynamic_search(self, train_loader, val_loader,
List of tuples (config, val_loss)
"""
+
+ # some preprocessing
+ for key, value in random_search_spaces.items():
+ if not isinstance(value, tuple):
+ random_search_spaces[key] = ([value], 'item')
+
config_count = 0
configs = []
- # stores tuples of (val_loss, config, model)
+ # stores tuples of (val_loss, config, model, config_id)
ranking = []
configs.append(self.random_search_spaces_to_config(random_search_spaces))
results = []
+ def update_ranking(ranking, val_loss, config, model, config_id):
+ ranking.append((val_loss, config, model, config_id))
+ ranking.sort(key=lambda x: x[0])
+ new_ranking = ranking[:ranks_considered]
+ dropped = ranking[ranks_considered:]
+
+ if model_path is not None:
+ for i in range(len(new_ranking)):
+ file_name = model_path + str(new_ranking[i][3]) + ".pt"
+ # check if the file already exists
+ if not os.path.isfile(file_name):
+ model.save_model(file_name)
+ for i in range(len(dropped)):
+ file_name = model_path + str(dropped[i][3]) + ".pt"
+ # delete the file if it exists
+ if os.path.isfile(file_name):
+ os.remove(file_name)
+ return new_ranking
+
+
def print_ranking(ranking):
+ print("\n" + "-"*30 + " RANKING " + "-"*30)
for i in range(len(ranking)):
- print("Rank {} with best loss {:.4f}: {}".format(i + 1, ranking[i][0], ranking[i][1]))
+ print("Config #{}: Rank {} with best loss {:.4f}:".format(ranking[i][3], i + 1, ranking[i][0]))
+ pprint.pprint(ranking[i][1])
+ print("-"*69 + "\n")
def adapt_search_space(ranking, random_search_space):
for name, (rng, mode) in random_search_space.items():
@@ -183,15 +255,18 @@ def adapt_search_space(ranking, random_search_space):
new_min = min(ranking, key=lambda x: x[1][name])[1][name]
new_max = max(ranking, key=lambda x: x[1][name])[1][name]
random_search_space[name] = ([new_min, new_max], mode)
- print("New search space:", random_search_space)
+ print("New search space:")
+ pprint.pprint(random_search_space)
return random_search_space
while config_count < num_search:
try:
config = configs[-1]
- print("\nEvaluating Config #{} [of {}]:\n".format(
- (config_count), num_search), config)
+ print("\nEvaluating Config #{} [of {}]:".format(
+ (config_count), num_search))
+ pprint.pprint(config)
+ print()
model = self.model_class(**config)
@@ -201,13 +276,13 @@ def adapt_search_space(ranking, random_search_space):
val_loss = trainer.train(optimizer, **config)
+ if evaluator is not None:
+ val_loss = evaluator.evaluate(model, val_loader)
+
# add into ranking
- if len(ranking) == 0:
- ranking.append((val_loss, config, model))
- else:
- ranking.append((val_loss, config, model))
- ranking.sort(key=lambda x: x[0])
- ranking = ranking[:ranks_considered]
+ ranking = update_ranking(ranking, val_loss, config, model, config_count)
+
+ results.append(val_loss)
print_ranking(ranking)
@@ -219,17 +294,25 @@ def adapt_search_space(ranking, random_search_space):
configs.append(self.random_search_spaces_to_config(random_search_spaces))
except KeyboardInterrupt:
break
+ except Exception as e:
+ print("During model training, got exception: ")
+ print(e)
best_val, best_model, \
best_config = ranking[0][0], ranking[0][2], ranking[0][1]
print("\nSearch done. Best Val Loss = {}".format(best_val))
- print("Best Config:", best_config)
+ print("Best Config #{}: ".format(ranking[0][3]))
+ pprint.pprint(best_config)
return best_model, best_config, list(zip(configs, results))
- def find_best_config(self, configs: List[Dict], train_loader: torch.utils.data.DataLoader,
- val_loader: torch.utils.data.DataLoader) -> Tuple[torch.nn.Module, Dict, List]:
+ def find_best_config(self,
+ configs: List[Dict],
+ train_loader: torch.utils.data.DataLoader,
+ val_loader: torch.utils.data.DataLoader,
+ evaluator: Evaluator = None
+ ) -> Tuple[torch.nn.Module, Dict, List]:
"""
Get a list of hyperparameter configs for random search or grid search,
trains a model on all configs and returns the one performing best
@@ -243,6 +326,9 @@ def find_best_config(self, configs: List[Dict], train_loader: torch.utils.data.D
The training data loader.
val_loader: torch.utils.data.DataLoader
The validation data loader.
+ evaluator: Evaluator
+ Evaluator class. Will be used instead of validation loss if given.
+ Its evaluate() should return a single metric.
Returns
-------
@@ -271,6 +357,9 @@ def find_best_config(self, configs: List[Dict], train_loader: torch.utils.data.D
val_loss = trainer.train(optimizer, **configs[i])
+ if evaluator is not None:
+ val_loss = evaluator.evaluate(model, val_loader)
+
results.append(val_loss)
if val_loss < best_val:
@@ -279,8 +368,10 @@ def find_best_config(self, configs: List[Dict], train_loader: torch.utils.data.D
best_model = model
- print("\nSearch done. Best Val Loss = {}".format(best_val))
- print("Best Config:", best_config)
+ print("\nSearch done. Best Metric = {}".format(best_val))
+ print("Best Config:")
+ pprint.pprint(best_config)
+
return best_model, best_config, list(zip(configs, results))