From d8dc3c7bd74fe03d4b50181ec6f045ceb677a0d0 Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 13:52:51 +0530 Subject: [PATCH 1/9] Create sidebar.py --- templates/Style Transfer_Pytorch/sidebar.py | 1 + 1 file changed, 1 insertion(+) create mode 100644 templates/Style Transfer_Pytorch/sidebar.py diff --git a/templates/Style Transfer_Pytorch/sidebar.py b/templates/Style Transfer_Pytorch/sidebar.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/templates/Style Transfer_Pytorch/sidebar.py @@ -0,0 +1 @@ + From f6f69bd1a318c6b8b47fcdc2ad677fa5a619c7f9 Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:05:42 +0530 Subject: [PATCH 2/9] Update sidebar.py --- templates/Style Transfer_Pytorch/sidebar.py | 92 +++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/templates/Style Transfer_Pytorch/sidebar.py b/templates/Style Transfer_Pytorch/sidebar.py index 8b13789..fa56e43 100644 --- a/templates/Style Transfer_Pytorch/sidebar.py +++ b/templates/Style Transfer_Pytorch/sidebar.py @@ -1 +1,93 @@ +import streamlit as st +MODELS = { + "VGG": { + "VGG11": "vgg11", + "VGG11 with batch normalization": "vgg11_bn", + "VGG13": "vgg13", + "VGG13 with batch normalization": "vgg13_bn", + "VGG16": "vgg16", + "VGG16 with batch normalization": "vgg16_bn", + "VGG19": "vgg19", + "VGG19 with batch normalization": "vgg19_bn", + }, +} + + +# Define possible optimizers in a dict. +# Format: optimizer -> default learning rate +OPTIMIZERS = { + "Adam": 0.001, + "Adadelta": 1.0, + "Adagrad": 0.01, + "Adamax": 0.002, + "RMSprop": 0.01, + "SGD": 0.1, +} + + + +def show(): + """Shows the sidebar components for the template and returns user inputs as dict.""" + + # `show()` is the only method required in this module. You can add any other code + # you like above or below. + + inputs = {} # dict to store all user inputs until return + + with st.sidebar: + + # Render all template-specific sidebar components here. + + # Use ## to denote sections. Common sections for training templates: + # Model, Input data, Preprocessing, Training, Visualizations + st.write("## Model") + + # Store all user inputs in the `inputs` dict. This will be passed to the code + # template later. + inputs["model"] = st.selectbox("Which model?", list(MODELS.keys())) + + + inputs["pretrained"] = st.checkbox("Use pre-trained model (Suggested Use is with a pretrained model)") + + st.write("## Input data") + inputs["data_format"] = st.selectbox( + "Which data do you want to use?", + ("Public Image File", "Custom Image files"), + ) + + if input["data_format"]== "Public Image File": + st.write(""" + ``` + It's a public image so you won't have to do anything. + ``` + """) + + elif input["data_format"]== "Custom Image File": + st.write(""" + ``` + Make sure you have two files the content image and the sytle image. + ``` + """) + + inputs["loss"] = st.selectbox( + "Loss function", ("Mean Squared Error Loss") + ) + + inputs["optimizer"] = st.selectbox("Optimizer", list(OPTIMIZERS.keys())) + + default_lr = OPTIMIZERS[inputs["optimizer"]] + inputs["lr"] = st.number_input( + "Learning rate", 0.000, None, default_lr, format="%f" + ) + + inputs["num_epochs"] = st.number_input("Epochs", 1, None, 5) + + + return inputs + + +# To test the sidebar independent of the app or template, just run +# `streamlit run sidebar.py` from within this folder. +if __name__ == "__main__": + show() From caf7050893b5c99a1892cef6ffe3139360adc231 Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:08:59 +0530 Subject: [PATCH 3/9] Create code-template.py.jinja --- .../code-template.py.jinja | 336 ++++++++++++++++++ 1 file changed, 336 insertions(+) create mode 100644 templates/Style Transfer_Pytorch/code-template.py.jinja diff --git a/templates/Style Transfer_Pytorch/code-template.py.jinja b/templates/Style Transfer_Pytorch/code-template.py.jinja new file mode 100644 index 0000000..2b7366a --- /dev/null +++ b/templates/Style Transfer_Pytorch/code-template.py.jinja @@ -0,0 +1,336 @@ +# Before running, install required packages: +{% if notebook %} + +! +{%- else %} +# +{%- endif %} + pip install numpy torch torchvision pytorch-ignite{% if visualization_tool == "Tensorboard" %} tensorboardX tensorboard{% endif %}{% if visualization_tool == "Weights & Biases" %} wandb{% endif %}{% if visualization_tool == "comet.ml" %} comet_ml{% endif %}{% if visualization_tool == "Aim" %} aim{% endif %} +{% if visualization_tool == "Weights & Biases" %} + + +# Also, you need to login to Weights & Biases on the terminal: +{% if notebook %} + +! wandb login +{% else %} +# wandb login +{% endif %} +{% endif %} +{% if notebook %} + + +# --- +{% endif %} + + +{% if visualization_tool == "comet.ml" %} +from comet_ml import Experiment # has to be 1st import +{% endif %} +import numpy as np +import torch +from torch import optim, nn +from torch.utils.data import DataLoader, TensorDataset +from torchvision import models, datasets, transforms +from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator +from ignite.metrics import Accuracy, Loss +{% if data_format == "Image files" %} +import urllib +import zipfile +{% endif %} +{% if visualization_tool == "Tensorboard" or checkpoint %} +from datetime import datetime +{% endif %} +{% if visualization_tool == "Tensorboard" %} +from tensorboardX import SummaryWriter +{% elif visualization_tool == "Aim" %} +from aim import Session +{% elif visualization_tool == "Weights & Biases" %} +import wandb +{% endif %} +{% if checkpoint %} +from pathlib import Path +{% endif %} + +{% if data_format == "Numpy arrays" %} +def fake_data(): + # 4 images of shape 1x16x16 with labels 0, 1, 2, 3 + return [np.random.rand(4, 1, 16, 16), np.arange(4)] + +{% elif data_format == "Image files" %} +# COMMENT THIS OUT IF YOU USE YOUR OWN DATA. +# Download example data into ./data/image-data (4 image files, 2 for "dog", 2 for "cat"). +url = "https://github.com/jrieke/traingenerator/raw/main/data/fake-image-data.zip" +zip_path, _ = urllib.request.urlretrieve(url) +with zipfile.ZipFile(zip_path, "r") as f: + f.extractall("data") + +{% endif %} + +{{ header("Setup") }} +{% if data_format == "Numpy arrays" %} +# INSERT YOUR DATA HERE +# Expected format: [images, labels] +# - images has array shape (num samples, color channels, height, width) +# - labels has array shape (num samples, ) +train_data = fake_data() # required +val_data = fake_data() # optional +test_data = None # optional +{% elif data_format == "Image files" %} +# INSERT YOUR DATA HERE +# Expected format: One folder per class, e.g. +# train +# --- dogs +# | +-- lassie.jpg +# | +-- komissar-rex.png +# --- cats +# | +-- garfield.png +# | +-- smelly-cat.png +# +# Example: https://github.com/jrieke/traingenerator/tree/main/data/image-data +train_data = "data/image-data" # required +val_data = "data/image-data" # optional +test_data = None # optional +{% elif data_format == "Public dataset"%} +# Dataset {{ dataset }} will be loaded further down. +{% endif %} + +# Set up hyperparameters. +lr = {{ lr }} +batch_size = {{ batch_size }} +num_epochs = {{ num_epochs }} + +# Set up logging. +{% if visualization_tool == "Tensorboard" or checkpoint %} +experiment_id = datetime.now().strftime('%Y-%m-%d_%H-%M-%S') +{% endif %} +{% if visualization_tool == "Tensorboard" %} +writer = SummaryWriter(logdir=f"logs/{experiment_id}") +{% elif visualization_tool == "Aim" %} +aim_session = Session({% if aim_experiment %}experiment="{{ aim_experiment }}"{% endif %}) +aim_session.set_params({"lr": lr, "batch_size": batch_size, "num_epochs": num_epochs}, name="hparams") +{% elif visualization_tool == "Weights & Biases" %} +wandb.init( +{% if wb_project %} + project="{{ wb_project }}", +{% endif %} +{% if wb_name %} + name="{{ wb_name }}", +{% endif %} + config={"lr": lr, "batch_size": batch_size, "num_epochs": num_epochs} +) +{% elif visualization_tool == "comet.ml" %} +experiment = Experiment("{{ comet_api_key }}"{% if comet_project %}, project_name="{{ comet_project }}"{% endif %}) +{% endif %} +{% if checkpoint %} +checkpoint_dir = Path(f"checkpoints/{experiment_id}") +checkpoint_dir.mkdir(parents=True, exist_ok=True) +{% endif %} +print_every = {{ print_every }} # batches + +# Set up device. +{% if gpu %} +use_cuda = torch.cuda.is_available() +{% else %} +use_cuda = False +{% endif %} +device = torch.device("cuda" if use_cuda else "cpu") + + +{% if data_format == "Public dataset" %} +{{ header("Dataset & Preprocessing") }} +def load_data(train): + # Download and transform dataset. + transform = transforms.Compose([ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + {% if dataset == "MNIST" or dataset == "FashionMNIST"%} + transforms.Lambda(lambda x: x.repeat(3, 1, 1)), # grayscale to RGB + {% endif %} + ]) + dataset = datasets.{{ dataset }}("./data", train=train, download=True, transform=transform) + + # Wrap in data loader. + if use_cuda: + kwargs = {"pin_memory": True, "num_workers": 1} + else: + kwargs = {} + loader = DataLoader(dataset, batch_size=batch_size, shuffle=train, **kwargs) + return loader + +train_loader = load_data(train=True) +val_loader = None +test_loader = load_data(train=False) +{% else %} +{{ header("Preprocessing") }} +def preprocess(data, name): + if data is None: # val/test can be empty + return None + + {% if data_format == "Image files" %} + # Read image files to pytorch dataset. + transform = transforms.Compose([ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + {# TODO: Maybe add normalization option even if model is not pretrained #} + {% if pretrained %} + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + {% endif %} + ]) + dataset = datasets.ImageFolder(data, transform=transform) + {% elif data_format == "Numpy arrays" %} + images, labels = data + + # Rescale images to 0-255 and convert to uint8. + # Note: This is done for each dataset individually, which is usually ok if all + # datasets look similar. If not, scale all datasets based on min/ptp of train set. + images = (images - np.min(images)) / np.ptp(images) * 255 + images = images.astype(np.uint8) + + # If images are grayscale, convert to RGB by duplicating channels. + if images.shape[1] == 1: + images = np.stack((images[:, 0],) * 3, axis=1) + + # Resize images and transform images torch tensor. + images = images.transpose((0, 2, 3, 1)) # channels-last, required for transforms.ToPILImage + transform = transforms.Compose([ + transforms.ToPILImage(), + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + {% if pretrained %} + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + {% endif %} + ]) + {# TODO: This is quite ugly and very inefficient #} + images = torch.stack(list(map(transform, images))) + + # Convert labels to tensors. + labels = torch.from_numpy(labels).long() + + # Construct dataset. + dataset = TensorDataset(images, labels) + {% endif %} + + # Wrap in data loader. + {% if gpu %} + if use_cuda: + kwargs = {"pin_memory": True, "num_workers": 1} + else: + kwargs = {} + loader = DataLoader(dataset, batch_size=batch_size, shuffle=(name=="train"), **kwargs) + {% else %} + loader = DataLoader(dataset, batch_size=batch_size, shuffle=(name=="train")) + {% endif %} + return loader + +train_loader = preprocess(train_data, "train") +val_loader = preprocess(val_data, "val") +test_loader = preprocess(test_data, "test") +{% endif %} + + +{{ header("Model") }} +# Set up model, loss, optimizer. +model = models.{{ model_func }}(pretrained={{ pretrained }}) +{# TODO: Maybe enable this by default, so that people can adapt num_classes afterward. #} +{% if num_classes != 1000 %} +num_classes = {{ num_classes }} +{% if "resnet" in model_func %} +model.fc = torch.nn.Linear(in_features=model.fc.in_features, out_features=num_classes, bias=True) +{% elif "alexnet" in model_func or "vgg" in model_func %} +model.classifier[-1] = torch.nn.Linear(in_features=model.classifier[-1].in_features, out_features=num_classes, bias=True) +{% elif "densenet" in model_func %} +model.classifier = torch.nn.Linear(in_features=model.classifier.in_features, out_features=num_classes, bias=True) +{% endif %} +{% endif %} +model = model.to(device) +loss_func = nn.{{ loss }}() +optimizer = optim.{{ optimizer }}(model.parameters(), lr=lr) + +{% if visualization_tool == "Weights & Biases" %} +# Log gradients and model parameters to W&B. +wandb.watch(model) + +{% endif %} + +{{ header("Training") }} +# Set up pytorch-ignite trainer and evaluator. +trainer = create_supervised_trainer( + model, + optimizer, + loss_func, + device=device, +) +{# TODO: Atm, the train metrics get accumulated, see torch_models.py #} +metrics = { + "accuracy": Accuracy(), + "loss": Loss(loss_func), +} +evaluator = create_supervised_evaluator( + model, metrics=metrics, device=device +) + +@trainer.on(Events.ITERATION_COMPLETED(every=print_every)) +def log_batch(trainer): + batch = (trainer.state.iteration - 1) % trainer.state.epoch_length + 1 + print( + f"Epoch {trainer.state.epoch} / {num_epochs}, " + f"batch {batch} / {trainer.state.epoch_length}: " + f"loss: {trainer.state.output:.3f}" + ) + +@trainer.on(Events.EPOCH_COMPLETED) +def log_epoch(trainer): + print(f"Epoch {trainer.state.epoch} / {num_epochs} average results: ") + + def log_results(name, metrics, epoch): + print( + f"{name + ':':6} loss: {metrics['loss']:.3f}, " + f"accuracy: {metrics['accuracy']:.3f}" + ) + {% if visualization_tool == "Tensorboard" %} + writer.add_scalar(f"{name}_loss", metrics["loss"], epoch) + writer.add_scalar(f"{name}_accuracy", metrics["accuracy"], epoch) + {% elif visualization_tool == "Aim" %} + aim_session.track(metrics["loss"], name="loss", subset=name, epoch=epoch) + aim_session.track(metrics["accuracy"], name="accuracy", subset=name, epoch=epoch) + {% elif visualization_tool == "Weights & Biases" %} + wandb.log({f"{name}_loss": metrics["loss"], f"{name}_accuracy": metrics["accuracy"]}) + {% elif visualization_tool == "comet.ml" %} + experiment.log_metric(f"{name}_loss", metrics["loss"]) + experiment.log_metric(f"{name}_accuracy", metrics["accuracy"]) + {% endif %} + + # Train data. + evaluator.run(train_loader) + log_results("train", evaluator.state.metrics, trainer.state.epoch) + + # Val data. + if val_loader: + evaluator.run(val_loader) + log_results("val", evaluator.state.metrics, trainer.state.epoch) + + # Test data. + if test_loader: + evaluator.run(test_loader) + log_results("test", evaluator.state.metrics, trainer.state.epoch) + + print() + print("-" * 80) + print() + +{# TODO: Maybe use this instead: https://pytorch.org/ignite/handlers.html#ignite.handlers.ModelCheckpoint #} +{% if checkpoint %} +@trainer.on(Events.EPOCH_COMPLETED) +def checkpoint_model(trainer): + torch.save(model, checkpoint_dir / f"model-epoch{trainer.state.epoch}.pt") + +{% endif %} +# Start training. +trainer.run(train_loader, max_epochs=num_epochs) +{% if visualization_tool == "Weights & Biases" %} +wandb.finish() +{% endif %} From 9a7bfd79fda1a15d4fa6dae5183cd54a22612fb0 Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:19:41 +0530 Subject: [PATCH 4/9] Update sidebar.py --- templates/Style Transfer_Pytorch/sidebar.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/templates/Style Transfer_Pytorch/sidebar.py b/templates/Style Transfer_Pytorch/sidebar.py index fa56e43..9d4ac08 100644 --- a/templates/Style Transfer_Pytorch/sidebar.py +++ b/templates/Style Transfer_Pytorch/sidebar.py @@ -1,16 +1,7 @@ import streamlit as st MODELS = { - "VGG": { - "VGG11": "vgg11", - "VGG11 with batch normalization": "vgg11_bn", - "VGG13": "vgg13", - "VGG13 with batch normalization": "vgg13_bn", - "VGG16": "vgg16", - "VGG16 with batch normalization": "vgg16_bn", - "VGG19": "vgg19", - "VGG19 with batch normalization": "vgg19_bn", - }, + "VGG": "VGG19": "vgg19", } From 73f5a19461a690598648827fb61f33c4c203e4f2 Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:21:15 +0530 Subject: [PATCH 5/9] Update sidebar.py --- templates/Style Transfer_Pytorch/sidebar.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/templates/Style Transfer_Pytorch/sidebar.py b/templates/Style Transfer_Pytorch/sidebar.py index 9d4ac08..19b98a0 100644 --- a/templates/Style Transfer_Pytorch/sidebar.py +++ b/templates/Style Transfer_Pytorch/sidebar.py @@ -38,7 +38,11 @@ def show(): # template later. inputs["model"] = st.selectbox("Which model?", list(MODELS.keys())) - + inputs["Device"] = st.selectbox( + "Which device would you like to train on?", + ("GPU", "CPU"), + ) + inputs["pretrained"] = st.checkbox("Use pre-trained model (Suggested Use is with a pretrained model)") st.write("## Input data") From f2342b502ee293214f06261518a5c5e0d9bdee5c Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:45:09 +0530 Subject: [PATCH 6/9] Update code-template.py.jinja --- .../code-template.py.jinja | 379 +++--------------- 1 file changed, 62 insertions(+), 317 deletions(-) diff --git a/templates/Style Transfer_Pytorch/code-template.py.jinja b/templates/Style Transfer_Pytorch/code-template.py.jinja index 2b7366a..06ff47c 100644 --- a/templates/Style Transfer_Pytorch/code-template.py.jinja +++ b/templates/Style Transfer_Pytorch/code-template.py.jinja @@ -1,336 +1,81 @@ -# Before running, install required packages: -{% if notebook %} - -! -{%- else %} -# -{%- endif %} - pip install numpy torch torchvision pytorch-ignite{% if visualization_tool == "Tensorboard" %} tensorboardX tensorboard{% endif %}{% if visualization_tool == "Weights & Biases" %} wandb{% endif %}{% if visualization_tool == "comet.ml" %} comet_ml{% endif %}{% if visualization_tool == "Aim" %} aim{% endif %} -{% if visualization_tool == "Weights & Biases" %} - - -# Also, you need to login to Weights & Biases on the terminal: -{% if notebook %} - -! wandb login -{% else %} -# wandb login -{% endif %} -{% endif %} -{% if notebook %} - - -# --- -{% endif %} - - -{% if visualization_tool == "comet.ml" %} -from comet_ml import Experiment # has to be 1st import -{% endif %} -import numpy as np import torch -from torch import optim, nn -from torch.utils.data import DataLoader, TensorDataset -from torchvision import models, datasets, transforms -from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator -from ignite.metrics import Accuracy, Loss -{% if data_format == "Image files" %} -import urllib -import zipfile -{% endif %} -{% if visualization_tool == "Tensorboard" or checkpoint %} -from datetime import datetime -{% endif %} -{% if visualization_tool == "Tensorboard" %} -from tensorboardX import SummaryWriter -{% elif visualization_tool == "Aim" %} -from aim import Session -{% elif visualization_tool == "Weights & Biases" %} -import wandb -{% endif %} -{% if checkpoint %} -from pathlib import Path -{% endif %} - -{% if data_format == "Numpy arrays" %} -def fake_data(): - # 4 images of shape 1x16x16 with labels 0, 1, 2, 3 - return [np.random.rand(4, 1, 16, 16), np.arange(4)] - -{% elif data_format == "Image files" %} -# COMMENT THIS OUT IF YOU USE YOUR OWN DATA. -# Download example data into ./data/image-data (4 image files, 2 for "dog", 2 for "cat"). -url = "https://github.com/jrieke/traingenerator/raw/main/data/fake-image-data.zip" -zip_path, _ = urllib.request.urlretrieve(url) -with zipfile.ZipFile(zip_path, "r") as f: - f.extractall("data") - -{% endif %} - -{{ header("Setup") }} -{% if data_format == "Numpy arrays" %} -# INSERT YOUR DATA HERE -# Expected format: [images, labels] -# - images has array shape (num samples, color channels, height, width) -# - labels has array shape (num samples, ) -train_data = fake_data() # required -val_data = fake_data() # optional -test_data = None # optional -{% elif data_format == "Image files" %} -# INSERT YOUR DATA HERE -# Expected format: One folder per class, e.g. -# train -# --- dogs -# | +-- lassie.jpg -# | +-- komissar-rex.png -# --- cats -# | +-- garfield.png -# | +-- smelly-cat.png -# -# Example: https://github.com/jrieke/traingenerator/tree/main/data/image-data -train_data = "data/image-data" # required -val_data = "data/image-data" # optional -test_data = None # optional -{% elif data_format == "Public dataset"%} -# Dataset {{ dataset }} will be loaded further down. -{% endif %} - -# Set up hyperparameters. -lr = {{ lr }} -batch_size = {{ batch_size }} -num_epochs = {{ num_epochs }} - -# Set up logging. -{% if visualization_tool == "Tensorboard" or checkpoint %} -experiment_id = datetime.now().strftime('%Y-%m-%d_%H-%M-%S') -{% endif %} -{% if visualization_tool == "Tensorboard" %} -writer = SummaryWriter(logdir=f"logs/{experiment_id}") -{% elif visualization_tool == "Aim" %} -aim_session = Session({% if aim_experiment %}experiment="{{ aim_experiment }}"{% endif %}) -aim_session.set_params({"lr": lr, "batch_size": batch_size, "num_epochs": num_epochs}, name="hparams") -{% elif visualization_tool == "Weights & Biases" %} -wandb.init( -{% if wb_project %} - project="{{ wb_project }}", -{% endif %} -{% if wb_name %} - name="{{ wb_name }}", -{% endif %} - config={"lr": lr, "batch_size": batch_size, "num_epochs": num_epochs} -) -{% elif visualization_tool == "comet.ml" %} -experiment = Experiment("{{ comet_api_key }}"{% if comet_project %}, project_name="{{ comet_project }}"{% endif %}) -{% endif %} -{% if checkpoint %} -checkpoint_dir = Path(f"checkpoints/{experiment_id}") -checkpoint_dir.mkdir(parents=True, exist_ok=True) -{% endif %} -print_every = {{ print_every }} # batches - -# Set up device. -{% if gpu %} -use_cuda = torch.cuda.is_available() -{% else %} -use_cuda = False -{% endif %} -device = torch.device("cuda" if use_cuda else "cpu") - - -{% if data_format == "Public dataset" %} -{{ header("Dataset & Preprocessing") }} -def load_data(train): - # Download and transform dataset. - transform = transforms.Compose([ - transforms.Resize(256), - transforms.CenterCrop(224), - transforms.ToTensor(), - {% if dataset == "MNIST" or dataset == "FashionMNIST"%} - transforms.Lambda(lambda x: x.repeat(3, 1, 1)), # grayscale to RGB - {% endif %} - ]) - dataset = datasets.{{ dataset }}("./data", train=train, download=True, transform=transform) - - # Wrap in data loader. - if use_cuda: - kwargs = {"pin_memory": True, "num_workers": 1} - else: - kwargs = {} - loader = DataLoader(dataset, batch_size=batch_size, shuffle=train, **kwargs) - return loader - -train_loader = load_data(train=True) -val_loader = None -test_loader = load_data(train=False) -{% else %} -{{ header("Preprocessing") }} -def preprocess(data, name): - if data is None: # val/test can be empty - return None - - {% if data_format == "Image files" %} - # Read image files to pytorch dataset. - transform = transforms.Compose([ - transforms.Resize(256), - transforms.CenterCrop(224), - transforms.ToTensor(), - {# TODO: Maybe add normalization option even if model is not pretrained #} - {% if pretrained %} - transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) - {% endif %} - ]) - dataset = datasets.ImageFolder(data, transform=transform) - {% elif data_format == "Numpy arrays" %} - images, labels = data - - # Rescale images to 0-255 and convert to uint8. - # Note: This is done for each dataset individually, which is usually ok if all - # datasets look similar. If not, scale all datasets based on min/ptp of train set. - images = (images - np.min(images)) / np.ptp(images) * 255 - images = images.astype(np.uint8) - - # If images are grayscale, convert to RGB by duplicating channels. - if images.shape[1] == 1: - images = np.stack((images[:, 0],) * 3, axis=1) +from torch import nn +from torch import optim +from torchvision import transforms +from PIL import Image +import torchvision.models as models +from torchvision.utils import save_image +import numpy as np - # Resize images and transform images torch tensor. - images = images.transpose((0, 2, 3, 1)) # channels-last, required for transforms.ToPILImage - transform = transforms.Compose([ - transforms.ToPILImage(), - transforms.Resize(256), - transforms.CenterCrop(224), - transforms.ToTensor(), - {% if pretrained %} - transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) - {% endif %} - ]) - {# TODO: This is quite ugly and very inefficient #} - images = torch.stack(list(map(transform, images))) +device = {{Device}} +model = models.vgg19(pretrained={{pretrained}}).features +print(model) - # Convert labels to tensors. - labels = torch.from_numpy(labels).long() +class VGG(nn.Module): + def __init__(self): + super(VGG,self).__init__() + self.choosen_features = ['0','5','10','19','28'] + self.model = models.vgg19(pretrained=True).features[:29] - # Construct dataset. - dataset = TensorDataset(images, labels) - {% endif %} + def forward(self,x): + features = [] - # Wrap in data loader. - {% if gpu %} - if use_cuda: - kwargs = {"pin_memory": True, "num_workers": 1} - else: - kwargs = {} - loader = DataLoader(dataset, batch_size=batch_size, shuffle=(name=="train"), **kwargs) - {% else %} - loader = DataLoader(dataset, batch_size=batch_size, shuffle=(name=="train")) - {% endif %} - return loader + for layer_num , layer in enumerate(self.model): + x = layer(x) -train_loader = preprocess(train_data, "train") -val_loader = preprocess(val_data, "val") -test_loader = preprocess(test_data, "test") -{% endif %} + if str(layer_num) in self.choosen_features: + features.append(x) + return features -{{ header("Model") }} -# Set up model, loss, optimizer. -model = models.{{ model_func }}(pretrained={{ pretrained }}) -{# TODO: Maybe enable this by default, so that people can adapt num_classes afterward. #} -{% if num_classes != 1000 %} -num_classes = {{ num_classes }} -{% if "resnet" in model_func %} -model.fc = torch.nn.Linear(in_features=model.fc.in_features, out_features=num_classes, bias=True) -{% elif "alexnet" in model_func or "vgg" in model_func %} -model.classifier[-1] = torch.nn.Linear(in_features=model.classifier[-1].in_features, out_features=num_classes, bias=True) -{% elif "densenet" in model_func %} -model.classifier = torch.nn.Linear(in_features=model.classifier.in_features, out_features=num_classes, bias=True) -{% endif %} -{% endif %} -model = model.to(device) -loss_func = nn.{{ loss }}() -optimizer = optim.{{ optimizer }}(model.parameters(), lr=lr) +model = VGG().to(device).eval() +loader = transforms.Compose([ + transforms.Resize((540,540)) , + transforms.ToTensor(), +]) +def load_image(image_name): + img = Image.open(image_name) + img = loader(img).unsqueeze(0) + return img.to(device) -{% if visualization_tool == "Weights & Biases" %} -# Log gradients and model parameters to W&B. -wandb.watch(model) +image_size = 540 +original_img = load_image('./style.jpg') +style_img = load_image('./style.jpg') -{% endif %} +generated = original_img.clone().requires_grad_(True) -{{ header("Training") }} -# Set up pytorch-ignite trainer and evaluator. -trainer = create_supervised_trainer( - model, - optimizer, - loss_func, - device=device, -) -{# TODO: Atm, the train metrics get accumulated, see torch_models.py #} -metrics = { - "accuracy": Accuracy(), - "loss": Loss(loss_func), -} -evaluator = create_supervised_evaluator( - model, metrics=metrics, device=device -) +total_steps = {{ epochs }} +learning_rate = {{ lr }} +alpha = 1 +beta = 0.01 +optimizer = optim.{{optimizer}}([generated],lr=learning_rate) -@trainer.on(Events.ITERATION_COMPLETED(every=print_every)) -def log_batch(trainer): - batch = (trainer.state.iteration - 1) % trainer.state.epoch_length + 1 - print( - f"Epoch {trainer.state.epoch} / {num_epochs}, " - f"batch {batch} / {trainer.state.epoch_length}: " - f"loss: {trainer.state.output:.3f}" - ) +for step in range(total_steps): + generated_features = model(generated) + original_img_features = model(original_img) + style_img_features = model(style_img) -@trainer.on(Events.EPOCH_COMPLETED) -def log_epoch(trainer): - print(f"Epoch {trainer.state.epoch} / {num_epochs} average results: ") + style_loss = original_loss = 0 - def log_results(name, metrics, epoch): - print( - f"{name + ':':6} loss: {metrics['loss']:.3f}, " - f"accuracy: {metrics['accuracy']:.3f}" - ) - {% if visualization_tool == "Tensorboard" %} - writer.add_scalar(f"{name}_loss", metrics["loss"], epoch) - writer.add_scalar(f"{name}_accuracy", metrics["accuracy"], epoch) - {% elif visualization_tool == "Aim" %} - aim_session.track(metrics["loss"], name="loss", subset=name, epoch=epoch) - aim_session.track(metrics["accuracy"], name="accuracy", subset=name, epoch=epoch) - {% elif visualization_tool == "Weights & Biases" %} - wandb.log({f"{name}_loss": metrics["loss"], f"{name}_accuracy": metrics["accuracy"]}) - {% elif visualization_tool == "comet.ml" %} - experiment.log_metric(f"{name}_loss", metrics["loss"]) - experiment.log_metric(f"{name}_accuracy", metrics["accuracy"]) - {% endif %} - - # Train data. - evaluator.run(train_loader) - log_results("train", evaluator.state.metrics, trainer.state.epoch) + for gen_fea,ori_fea,style_fea in zip(generated_features,original_img_features,style_img_features): + batch_size , channel,height,width = gen_fea.shape + original_loss += torch.mean((gen_fea - ori_fea) ** 2) - # Val data. - if val_loader: - evaluator.run(val_loader) - log_results("val", evaluator.state.metrics, trainer.state.epoch) + G = gen_fea.view(channel,height*width).mm( + gen_fea.view(channel,height*width).t() + ) - # Test data. - if test_loader: - evaluator.run(test_loader) - log_results("test", evaluator.state.metrics, trainer.state.epoch) + A = style_fea.view(channel,height*width).mm( + style_fea.view(channel,height*width).t() + ) - print() - print("-" * 80) - print() + style_loss += torch.mean((G - A)**2) -{# TODO: Maybe use this instead: https://pytorch.org/ignite/handlers.html#ignite.handlers.ModelCheckpoint #} -{% if checkpoint %} -@trainer.on(Events.EPOCH_COMPLETED) -def checkpoint_model(trainer): - torch.save(model, checkpoint_dir / f"model-epoch{trainer.state.epoch}.pt") + total_loss = alpha*original_loss + beta * style_loss + optimizer.zero_grad() + total_loss.backward() + optimizer.step() -{% endif %} -# Start training. -trainer.run(train_loader, max_epochs=num_epochs) -{% if visualization_tool == "Weights & Biases" %} -wandb.finish() -{% endif %} + if step % {{save}} == 0: + print(f'Total_Loss: {total_loss}') + save_image(generated,'Generated.png') From 8820d6ee66ad19b660510d3f4df5af7d5b68f62b Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:47:38 +0530 Subject: [PATCH 7/9] Update sidebar.py --- templates/Style Transfer_Pytorch/sidebar.py | 20 +++++++------------- 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/templates/Style Transfer_Pytorch/sidebar.py b/templates/Style Transfer_Pytorch/sidebar.py index 19b98a0..7a46a76 100644 --- a/templates/Style Transfer_Pytorch/sidebar.py +++ b/templates/Style Transfer_Pytorch/sidebar.py @@ -1,7 +1,7 @@ import streamlit as st MODELS = { - "VGG": "VGG19": "vgg19", + "VGG": "vgg19", } @@ -48,23 +48,17 @@ def show(): st.write("## Input data") inputs["data_format"] = st.selectbox( "Which data do you want to use?", - ("Public Image File", "Custom Image files"), + ("Custom Image files"), ) - if input["data_format"]== "Public Image File": + if input["data_format"]== "Custom Image files": st.write(""" ``` - It's a public image so you won't have to do anything. + Make sure you have style.jpg and content.jpg . ``` """) - - elif input["data_format"]== "Custom Image File": - st.write(""" - ``` - Make sure you have two files the content image and the sytle image. - ``` - """) - + + inputs["loss"] = st.selectbox( "Loss function", ("Mean Squared Error Loss") ) @@ -76,7 +70,7 @@ def show(): "Learning rate", 0.000, None, default_lr, format="%f" ) - inputs["num_epochs"] = st.number_input("Epochs", 1, None, 5) + inputs["num_epochs"] = st.number_input("Epochs", 1, None, 5000) return inputs From 242f7acb36dc7e9009aaa84a27020966fde841c1 Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:49:04 +0530 Subject: [PATCH 8/9] Update code-template.py.jinja --- templates/Style Transfer_Pytorch/code-template.py.jinja | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/templates/Style Transfer_Pytorch/code-template.py.jinja b/templates/Style Transfer_Pytorch/code-template.py.jinja index 06ff47c..a45c541 100644 --- a/templates/Style Transfer_Pytorch/code-template.py.jinja +++ b/templates/Style Transfer_Pytorch/code-template.py.jinja @@ -39,7 +39,7 @@ def load_image(image_name): return img.to(device) image_size = 540 -original_img = load_image('./style.jpg') +original_img = load_image('./content.jpg') style_img = load_image('./style.jpg') generated = original_img.clone().requires_grad_(True) @@ -76,6 +76,6 @@ for step in range(total_steps): total_loss.backward() optimizer.step() - if step % {{save}} == 0: + if step % {{ visualize_per_epoch }} == 0: print(f'Total_Loss: {total_loss}') save_image(generated,'Generated.png') From 0411903c14d87cfccad63fd32a27ccb8ce9df684 Mon Sep 17 00:00:00 2001 From: The-ML-Hero <67855775+The-ML-Hero@users.noreply.github.com> Date: Fri, 29 Jan 2021 14:49:23 +0530 Subject: [PATCH 9/9] Update sidebar.py --- templates/Style Transfer_Pytorch/sidebar.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/templates/Style Transfer_Pytorch/sidebar.py b/templates/Style Transfer_Pytorch/sidebar.py index 7a46a76..674d559 100644 --- a/templates/Style Transfer_Pytorch/sidebar.py +++ b/templates/Style Transfer_Pytorch/sidebar.py @@ -60,7 +60,7 @@ def show(): inputs["loss"] = st.selectbox( - "Loss function", ("Mean Squared Error Loss") + "Loss function", ("Style Gram Loss") ) inputs["optimizer"] = st.selectbox("Optimizer", list(OPTIMIZERS.keys())) @@ -71,6 +71,8 @@ def show(): ) inputs["num_epochs"] = st.number_input("Epochs", 1, None, 5000) + + inputs["visualize_per_epoch"] = st.number_input("Epochs", 1, None, 200) return inputs