Getting Started

Datasets

AshPy supports tf.data.Dataset format.

We highly encourage you to use Tensorflow Datasets to manage and use your datasets in an handy way.

pip install tfds-nightly

Classification

In order to create a dataset for classification:

import tensorflow_datasets as tfds

from ashpy.trainers import ClassifierTrainer

def extract_fn(example):
    return example["image"], example["label"]

def main():
    ds_train, ds_validation = tfds.load(name="mnist", split=["train", "validation"])

    # build the input pipeline
    ds_train = ds_train.batch(BATCH_SIZE).prefetch(1)
    ds_train = ds_train.map(extract_fn)

    # same for validation
    ...

    # define model, loss, optimizer
    ...

    # define the classifier trainer
    trainer = ClassifierTrainer(model, optimizer, loss, epochs, metrics, logdir=logdir)

    # train
    trainer.train(ds_train, ds_validation)

GANs

In order to create a datasets for a (Conditional) GANs:

import tensorflow_datasets as tfds

from ashpy.trainers import AdversarialTrainer

def extract_fn(example):
    # the ashpy input must be (real, condition), condition
    return (example["image"], example["label"]), example["label"]

def main():
    ds_train = tfds.load(name="mnist", split="train")

    # build the input pipeline
    ds_train = ds_train.batch(BATCH_SIZE).prefetch(1)
    ds_train = ds_train.map(extract_fn)

    # define models, losses, optimizers
    ...

    # define the adversarial trainer
    trainer = AdversarialTrainer(generator,
        discriminator,
        generator_optimizer,
        discriminator_optimizer,
        generator_loss,
        discriminator_loss,
        epochs,
        metrics,
        logdir,
    )

    # train
    trainer.train(ds_train)

Models

AshPy supports Keras models as inputs. You can use an AshPy predefined model or you can implement your own model.

Using an AshPy model

import tensorflow_datasets as tfds

from ashpy.trainers import ClassifierTrainer
from ashpy.models import UNet

def main():

    # create the dataset and the input pipeline

    # define models, loss, optimizer
    model = UNet(
        input_res,
        min_res,
        kernel_size,
        initial_filters,
        filters_cap,
        channels,
        use_dropout_encoder,
        use_dropout_decoder,
        dropout_prob,
        use_attention,
    )

    # define the classifier trainer
    trainer = AdversarialTrainer(generator,
        discriminator,
        generator_optimizer,
        discriminator_optimizer,
        generator_loss,
        discriminator_loss,
        epochs,
        metrics,
        logdir,
    )

    # train
    trainer.train(ds_train)

Creating a Model

It’s very easy to create a simple model, since AshPy’s models are Keras’ models.

from ashpy.layers import Attention, InstanceNormalization

def downsample(
    filters,
    apply_normalization=True,
    attention=False,
    activation=tf.keras.layers.LeakyReLU(alpha=0.2),
    size=3,
):
initializer = tf.random_normal_initializer(0.0, 0.02)

result = tf.keras.Sequential()
result.add(
    tf.keras.layers.Conv2D(
        filters,
        size,
        strides=2,
        padding="same",
        kernel_initializer=initializer,
        use_bias=not apply_normalization,
    )
)

if apply_normalization:
    result.add(InstanceNormalization())

result.add(activation)

if attention:
    result.add(Attention(filters))

return result


def upsample(
    filters,
    apply_dropout=False,
    apply_normalization=True,
    attention=False,
    activation=tf.keras.layers.ReLU(),
    size=3,
):
    initializer = tf.random_normal_initializer(0.0, 0.02)

    result = tf.keras.Sequential()
    result.add(tf.keras.layers.UpSampling2D(size=(2, 2)))
    result.add(tf.keras.layers.ZeroPadding2D(padding=(1, 1)))

    result.add(
        tf.keras.layers.Conv2D(
            filters,
            size,
            strides=1,
            padding="valid",
            kernel_initializer=initializer,
            use_bias=False,
        )
    )

    if apply_dropout:
        result.add(tf.keras.layers.Dropout(0.5))

    if apply_normalization:
        result.add(Normalizer())

    result.add(activation)

    if attention:
        result.add(Attention(filters))

    return result


def Generator(attention, output_channels=3):
    down_stack = [
        downsample(32, apply_normalization=False),  # 256
        downsample(32),  # 128
        downsample(64, attention=attention),  # 64
        downsample(64),  # 32
        downsample(64),  # 16
        downsample(128),  # 8
        downsample(128),  # 4
        downsample(256),  # 2
        downsample(512, apply_normalization=False),  # 1
    ]

    up_stack = [
        upsample(256, apply_dropout=True),  # 2
        upsample(128, apply_dropout=True),  # 4
        upsample(128, apply_dropout=True),  # 8
        upsample(64),  # 16
        upsample(64),  # 32
        upsample(64, attention=attention),  # 64
        upsample(32),  # 128
        upsample(32),  # 256
        upsample(32),  # 512
    ]

    inputs = tf.keras.layers.Input(shape=[None, None, 1])
    x = inputs

    # Downsampling through the model
    skips = []
    for down in down_stack:
        x = down(x)
        skips.append(x)

    skips = reversed(skips[:-1])

    # Upsampling and establishing the skip connections
    for up, skip in zip(up_stack, skips):
        x = up(x)
        x = tf.keras.layers.Concatenate()([x, skip])

    last = upsample(
        output_channels,
        activation=tf.keras.layers.Activation(tf.nn.tanh),
        apply_normalization=False,
    )

    x = last(x)

    return tf.keras.Model(inputs=inputs, outputs=x)

In this way we have created a new model to be used inside AshPy.

Inheriting from ashpy.models.Conv2DInterface

The third possibility you have to create a new model is to inherit from the ashpy.models.convolutional.interfaces.Conv2DInterface.

This class offers the basic methods to implement in a simple way a new model.

Creating a new Trainer

AshPy has different generics trainers. Trainers implement the basic training loop together with distribution strategy management and logging. By now the only distribution strategy handled is the tf.distribute.MirroredStrategy.

Complete Examples

Classifier

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
# Copyright 2019 Zuru Tech HK Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Example of Multi-GPU classifier trainer."""

import operator

import tensorflow as tf

from ashpy.losses import ClassifierLoss
from ashpy.metrics import ClassifierMetric
from ashpy.trainers import ClassifierTrainer


def main():
    """
    Train a multi-GPU classifier.

    How to use ash to training_set a classifier, measure the
    performance and perform model selection.
    """
    strategy = tf.distribute.MirroredStrategy()
    with strategy.scope():
        training_set, validation_set = tf.keras.datasets.mnist.load_data()

        def process(images, labels):
            data_images = tf.data.Dataset.from_tensor_slices((images)).map(
                lambda x: tf.reshape(x, (28 * 28,))
            )
            data_images = data_images.map(
                lambda x: tf.image.convert_image_dtype(x, tf.float32)
            )
            data_labels = tf.data.Dataset.from_tensor_slices((labels))
            dataset = tf.data.Dataset.zip((data_images, data_labels))
            dataset = dataset.batch(1024 * 1)
            return dataset

        training_set, validation_set = (
            process(training_set[0], training_set[1]),
            process(validation_set[0], validation_set[1]),
        )

        model = tf.keras.Sequential(
            [
                tf.keras.layers.Dense(10, activation=tf.nn.sigmoid),
                tf.keras.layers.Dense(10),
            ]
        )
        optimizer = tf.optimizers.Adam(1e-3)
        loss = ClassifierLoss(tf.losses.SparseCategoricalCrossentropy(from_logits=True))
        logdir = "testlog"
        epochs = 10

        metrics = [
            ClassifierMetric(
                tf.metrics.Accuracy(), model_selection_operator=operator.gt
            ),
            ClassifierMetric(
                tf.metrics.BinaryAccuracy(), model_selection_operator=operator.gt
            ),
        ]

        trainer = ClassifierTrainer(
            model=model,
            optimizer=optimizer,
            loss=loss,
            epochs=epochs,
            metrics=metrics,
            logdir=logdir,
        )
        trainer(training_set, validation_set)


if __name__ == "__main__":
    main()

GANs

BiGAN

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
# Copyright 2019 Zuru Tech HK Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Bigan dummy implementation."""

import operator

import tensorflow as tf
from tensorflow import keras

from ashpy.losses import DiscriminatorMinMax, EncoderBCE, GeneratorBCE
from ashpy.metrics import EncodingAccuracy
from ashpy.trainers import EncoderTrainer


def main():
    """Define the trainer and the models."""

    def real_gen():
        """Define generator of real values."""
        for _ in tf.range(100):
            yield ((10.0,), (0,))

    num_classes = 1
    latent_dim = 100

    generator = keras.Sequential([keras.layers.Dense(1)])

    left_input = tf.keras.layers.Input(shape=(1,))
    left = tf.keras.layers.Dense(10, activation=tf.nn.elu)(left_input)

    right_input = tf.keras.layers.Input(shape=(latent_dim,))
    right = tf.keras.layers.Dense(10, activation=tf.nn.elu)(right_input)

    net = tf.keras.layers.Concatenate()([left, right])
    out = tf.keras.layers.Dense(1)(net)

    discriminator = tf.keras.Model(inputs=[left_input, right_input], outputs=[out])

    encoder = keras.Sequential([keras.layers.Dense(latent_dim)])
    generator_bce = GeneratorBCE()
    encoder_bce = EncoderBCE()
    minmax = DiscriminatorMinMax()

    epochs = 100
    logdir = "log/adversarial/encoder"

    # Fake pre-trained classifier
    classifier = tf.keras.Sequential(
        [tf.keras.layers.Dense(10), tf.keras.layers.Dense(num_classes)]
    )

    metrics = [
        EncodingAccuracy(
            classifier, model_selection_operator=operator.gt, logdir=logdir
        )
    ]

    trainer = EncoderTrainer(
        generator=generator,
        discriminator=discriminator,
        encoder=encoder,
        generator_optimizer=tf.optimizers.Adam(1e-4),
        discriminator_optimizer=tf.optimizers.Adam(1e-5),
        encoder_optimizer=tf.optimizers.Adam(1e-6),
        generator_loss=generator_bce,
        discriminator_loss=minmax,
        encoder_loss=encoder_bce,
        epochs=epochs,
        metrics=metrics,
        logdir=logdir,
    )

    batch_size = 10
    discriminator_input = tf.data.Dataset.from_generator(
        real_gen, (tf.float32, tf.int64), ((1), (1))
    ).batch(batch_size)

    dataset = discriminator_input.map(
        lambda x, y: ((x, y), tf.random.normal(shape=(batch_size, latent_dim)))
    )

    trainer(dataset)


if __name__ == "__main__":
    main()

MNIST

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
# Copyright 2019 Zuru Tech HK Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Adversarial trainer example."""

import tensorflow as tf
from tensorflow import keras  # pylint: disable=no-name-in-module

from ashpy.losses import DiscriminatorMinMax, GeneratorBCE
from ashpy.metrics import InceptionScore
from ashpy.models.gans import ConvDiscriminator, ConvGenerator
from ashpy.trainers import AdversarialTrainer


def main():
    """Adversarial trainer example."""
    strategy = tf.distribute.MirroredStrategy()
    with strategy.scope():

        generator = ConvGenerator(
            layer_spec_input_res=(7, 7),
            layer_spec_target_res=(28, 28),
            kernel_size=(5, 5),
            initial_filters=256,
            filters_cap=16,
            channels=1,
        )

        discriminator = ConvDiscriminator(
            layer_spec_input_res=(28, 28),
            layer_spec_target_res=(7, 7),
            kernel_size=(5, 5),
            initial_filters=32,
            filters_cap=128,
            output_shape=1,
        )

        # Losses
        generator_bce = GeneratorBCE()
        minmax = DiscriminatorMinMax()

        # Trainer
        logdir = "log/adversarial"

        # InceptionScore: keep commented until the issues
        # https://github.com/tensorflow/tensorflow/issues/28599
        # https://github.com/tensorflow/hub/issues/295
        # Haven't been solved and merged into tf2

        metrics = [
            # InceptionScore(
            #    InceptionScore.get_or_train_inception(
            #        mnist_dataset,
            #        "mnist",
            #        num_classes=10,
            #        epochs=1,
            #        fine_tuning=False,
            #        logdir=logdir,
            #    ),
            #    model_selection_operator=operator.gt,
            #    logdir=logdir,
            # )
        ]

        epochs = 50
        trainer = AdversarialTrainer(
            generator=generator,
            discriminator=discriminator,
            generator_optimizer=tf.optimizers.Adam(1e-4),
            discriminator_optimizer=tf.optimizers.Adam(1e-4),
            generator_loss=generator_bce,
            discriminator_loss=minmax,
            epochs=epochs,
            metrics=metrics,
            logdir=logdir,
        )

        batch_size = 512

        # Real data
        mnist_x, mnist_y = keras.datasets.mnist.load_data()[0]

        def iterator():
            """Define an iterator in order to do not load in memory all the dataset."""
            for image, label in zip(mnist_x, mnist_y):
                yield tf.image.convert_image_dtype(
                    tf.expand_dims(image, -1), tf.float32
                ), tf.expand_dims(label, -1)

        real_data = (
            tf.data.Dataset.from_generator(
                iterator, (tf.float32, tf.int64), ((28, 28, 1), (1,))
            )
            .batch(batch_size)
            .prefetch(1)
        )

        # Add noise in the same dataset, just by mapping.
        # The return type of the dataset must be: tuple(tuple(a,b), noise)
        dataset = real_data.map(
            lambda x, y: ((x, y), tf.random.normal(shape=(batch_size, 100)))
        )

        trainer(dataset)


if __name__ == "__main__":
    main()

Facades (Pix2Pix)

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
# Copyright 2019 Zuru Tech HK Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Pix2Pix on Facades Datasets dummy implementation.

Input Pipeline taken from: https://www.tensorflow.org/beta/tutorials/generative/pix2pix
"""
from pathlib import Path

import tensorflow as tf

from ashpy import LogEvalMode
from ashpy.losses.gan import (
    AdversarialLossType,
    Pix2PixLoss,
    get_adversarial_loss_discriminator,
)
from ashpy.models.convolutional.discriminators import PatchDiscriminator
from ashpy.models.convolutional.unet import FUNet
from ashpy.trainers.gan import AdversarialTrainer

_URL = "https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz"

PATH_TO_ZIP = tf.keras.utils.get_file("facades.tar.gz", origin=_URL, extract=True)
PATH = Path(PATH_TO_ZIP).parent / "facades"

BUFFER_SIZE = 100
BATCH_SIZE = 1
IMG_WIDTH = 256
IMG_HEIGHT = 256


def load(image_file):
    """Load the image from file path."""
    image = tf.io.read_file(image_file)
    image = tf.image.decode_jpeg(image)

    width = tf.shape(image)[1]

    width = width // 2
    real_image = image[:, :width, :]
    input_image = image[:, width:, :]

    input_image = tf.cast(input_image, tf.float32)
    real_image = tf.cast(real_image, tf.float32)

    return input_image, real_image


def resize(input_image, real_image, height, width):
    """Resize input_image and real_image to height x width."""
    input_image = tf.image.resize(
        input_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR
    )
    real_image = tf.image.resize(
        real_image, [height, width], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR
    )

    return input_image, real_image


def random_crop(input_image, real_image):
    """Random crop both input_image and real_image."""
    stacked_image = tf.stack([input_image, real_image], axis=0)
    cropped_image = tf.image.random_crop(
        stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3]
    )

    return cropped_image[0], cropped_image[1]


def normalize(input_image, real_image):
    """Normalize images in [-1, 1]."""
    input_image = (input_image / 127.5) - 1
    real_image = (real_image / 127.5) - 1

    return input_image, real_image


def load_image_train(image_file):
    """Load and process the image_file to be ready for the training."""
    input_image, real_image = load(image_file)
    input_image, real_image = random_jitter(input_image, real_image)
    input_image, real_image = normalize(input_image, real_image)

    return input_image, real_image


@tf.function
def random_jitter(input_image, real_image):
    """Apply random jitter to both input_image and real_image."""
    # resizing to 286 x 286 x 3
    input_image, real_image = resize(input_image, real_image, 286, 286)

    # randomly cropping to 256 x 256 x 3
    input_image, real_image = random_crop(input_image, real_image)

    if tf.random.uniform(()) > 0.5:
        # random mirroring
        input_image = tf.image.flip_left_right(input_image)
        real_image = tf.image.flip_left_right(real_image)

    return input_image, real_image


def main(
    kernel_size=5,
    learning_rate_d=2e-4,
    learning_rate_g=2e-4,
    g_input_res=IMG_WIDTH,
    g_min_res=1,
    g_initial_filters=64,
    g_filters_cap=512,
    use_dropout_encoder=False,
    use_dropout_decoder=True,
    d_target_res=32,
    d_initial_filters=64,
    d_filters_cap=512,
    use_dropout_discriminator=False,
    dataset_name="facades",
    resolution=256,
    epochs=100_000,
    dropout_prob=0.3,
    l1_loss_weight=100,
    gan_loss_weight=1,
    use_attention_d=False,
    use_attention_g=False,
    channels=3,
    gan_loss_type=AdversarialLossType.LSGAN,
):
    """Define Trainer and models."""
    generator = FUNet(
        input_res=g_input_res,
        min_res=g_min_res,
        kernel_size=kernel_size,
        initial_filters=g_initial_filters,
        filters_cap=g_filters_cap,
        channels=channels,  # color_to_label_tensor.shape[0],
        use_dropout_encoder=use_dropout_encoder,
        use_dropout_decoder=use_dropout_decoder,
        dropout_prob=dropout_prob,
        use_attention=use_attention_g,
    )
    discriminator = PatchDiscriminator(
        input_res=resolution,
        min_res=d_target_res,
        initial_filters=d_initial_filters,
        kernel_size=kernel_size,
        filters_cap=d_filters_cap,
        use_dropout=use_dropout_discriminator,
        dropout_prob=dropout_prob,
        use_attention=use_attention_d,
    )

    discriminator_loss = get_adversarial_loss_discriminator(gan_loss_type)()
    generator_loss = Pix2PixLoss(
        l1_loss_weight=l1_loss_weight,
        adversarial_loss_weight=gan_loss_weight,
        adversarial_loss_type=gan_loss_type,
    )

    metrics = []
    logdir = Path("log") / dataset_name / "run2"

    if not logdir.exists():
        logdir.mkdir(parents=True)

    trainer = AdversarialTrainer(
        generator=generator,
        discriminator=discriminator,
        generator_optimizer=tf.optimizers.Adam(learning_rate_g, beta_1=0.5),
        discriminator_optimizer=tf.optimizers.Adam(learning_rate_d, beta_1=0.5),
        generator_loss=generator_loss,
        discriminator_loss=discriminator_loss,
        epochs=epochs,
        metrics=metrics,
        logdir=logdir,
        log_eval_mode=LogEvalMode.TEST,
    )

    train_dataset = tf.data.Dataset.list_files(PATH + "train/*.jpg")
    train_dataset = train_dataset.shuffle(BUFFER_SIZE)
    train_dataset = train_dataset.map(load_image_train)
    train_dataset = train_dataset.batch(BATCH_SIZE)

    train_dataset = train_dataset.map(lambda x, y: ((y, x), x))

    trainer(
        # generator_input,
        train_dataset
    )


if __name__ == "__main__":
    main()