python - How Implement on the fly augmentation in the tensorflow?

Question

I want implement augmentation in a 3D dataset with an tensorflow model.

the augmentation function is like this:

    def augmentation(img, label):
            
            p = .5
            print('augmentation')
            
            if random.random() > p:
                img = tf.numpy_function(augment_noise, [img], tf.double)
                
            if random.random() > p:
                img = tf.numpy_function(flip_x, [img], tf.double)
    
            if random.random() > p:
                img = tf.numpy_function(augment_scale, [img], tf.double)
    
            if random.random() > p:
                img = tf.numpy_function(distort_elastic_cv2, [img], tf.double)    
            
           
            img = tf.image.convert_image_dtype(img, tf.float32)
            
            return img, label

The augment functions is not implemented in tensorflow.

Tensorflow code for using this function is as follow:

ds_train = tf.data.Dataset.from_tensor_slices((image_train, label_train))
ds_valid = tf.data.Dataset.from_tensor_slices((image_val, label_val))


batch_size = 16
repeat_count = int((1000 * batch_size)/len(image_train))
# AUTOTUNE =  tf.data.experimental.AUTOTUNE # tf.data.AUTOTUNE
AUTOTUNE = 16

# Augment the on the fly during training.
ds_train = (
    ds_train.shuffle(len(ds_train)).repeat(repeat_count)
    .map(augmentation, num_parallel_calls=AUTOTUNE)
    .batch(batch_size)
    .prefetch(buffer_size=AUTOTUNE)
)


ds_valid = (
    ds_valid.batch(batch_size)
    .prefetch(buffer_size=AUTOTUNE)
)

initial_epoch = 0
epochs = 1000
H = model.fit(ds_train, validation_data=ds_valid,initial_epoch=initial_epoch,
               epochs = epochs,
              callbacks = chkpts, use_multiprocessing=False, workers=1, verbose=2)

I want in each epoch randomly select about 1000 batch from the dataset, and then augment them on the. The I calculate repeat_count to create 1000 batch with size of batch_size.

The problem is that I don't know is model call augmentation function in each epoch and implies it to each images of the batch(I mean 161000 images in each epoch), so I added print in the augmentation function, that it only printed one times, not in each epoch, or for each image. Is the augmentation function call 161000 times in each epoch?

Also the Utilization of cpu and gpu is vary in each run of the code. some time the Utilization of cpu is about 25% and gpu is 30, but in almost of runs it is 100% and 5.

How can solve the two problem?

question from:https://stackoverflow.com/questions/66063619/how-implement-on-the-fly-augmentation-in-the-tensorflow

Answer 1

Your string is printed once because it's calling once to make a Tensorflow graph. If you use tf.print to print, it will be part of the graph so it will print every time.

Copy/paste this:

import tensorflow as tf
import matplotlib.pyplot as plt
from sklearn.datasets import load_sample_image
import numpy as np
import random

imgs = np.stack([load_sample_image('flower.jpg') for i in range(4*4)], axis=0)

def augmentation(img):      
        p = .5
        tf.print('augmentation successful!')  
        img = tf.image.convert_image_dtype(img, tf.float32)
        return img 
    

ds_train = tf.data.Dataset.from_tensor_slices(imgs)


batch_size = 16
repeat_count = 10
AUTOTUNE = 16

ds_train = (
    ds_train.shuffle(len(ds_train)).repeat(repeat_count)
    .map(augmentation, num_parallel_calls=AUTOTUNE)
    .batch(batch_size)
    .prefetch(buffer_size=AUTOTUNE)
)

for i in ds_train:
    pass

augmentation successful!
augmentation successful!
augmentation successful!
augmentation successful!
augmentation successful!
augmentation successful!
augmentation successful!
augmentation successful!
augmentation successful!