adv_training_experiments.py

import keras.backend
import numpy as np

import model_functions
import utility_functions
from model_functions import *
from callbacks import *
from training_techniques import *
import socket
import segmentor_training
import denoisor_training
import classifier_training
import enum
from PIL import ImageFile
import tensorflow as tf

import gc

###########################
### Setup Logging Level ###
###########################

class LOGGING_LEVEL(enum.Enum):
    DEBUG = 10
    INFO = 20
    WARNING = 30
    ERROR = 40
    FATAL = 50


tf_level = LOGGING_LEVEL.ERROR
tf_level = tf_level.value

###################################
### Setup Environment Variables ###
###################################
os.environ['GCS_READ_CACHE_DISABLED'] = '1'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = str(tf_level)
os.environ['AUTOGRAPH_VERBOSITY'] = str(tf_level)

os.environ['TF_DEVICE_MIN_SYS_MEMORY_IN_MB'] = '256'
os.environ['TF_ENABLE_GPU_GARBAGE_COLLECTION'] = 'false'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'
os.environ['CUDA_LAUNCH_BLOCKING'] = '0'

tf.get_logger().setLevel(tf_level)
tf.autograph.set_verbosity(tf_level)

ImageFile.LOAD_TRUNCATED_IMAGES = True
np.random.seed(0)

try:
    import pydevd

    DEBUGGING = True
except ImportError:
    DEBUGGING = False

if DEBUGGING:
    tf.config.run_functions_eagerly(True)
    # tf.data.experimental.enable_debug_mode()
    tf.autograph.experimental.do_not_convert()
else:
    tf.config.run_functions_eagerly(False)

print(tf.config.list_physical_devices())

input_shape = (48, 48)
print("input shape", str(input_shape))
output_shape = input_shape
num_classes = num_classes()

datasets = "datasets/"
unlabelled_train_path = datasets + "denoising/train"
unlabelled_test_path = datasets + "denoising/test"

seg_train_path = datasets + "sim/train"
seg_test_path = datasets + "sim/test"

seg_epochs = 40
repeats = 3

discr_weights = [0.05, 0.01]
discr_depths = [1, 4]
def freeze(layers):
    for layer in layers:
        layer.trainable = False


def sub_path(discr_weight, i):
    return "w_" + str(discr_weight) + "_" + str(i)


# _, _, X, y = get_segmentation_data(seg_train_path, True, input_shape, 0, include_original_data=True)
# _, _, X_test, y_test = get_segmentation_data(seg_test_path, True, input_shape, 0, include_original_data=True)

choice = "cs"
betas = [1, 2, 3]

datasets = "datasets/"

train_path = datasets + "sim/train"
test_path = datasets + "sim/test"

y_train_path = train_path + "/labels"
y_test_path = test_path + "/labels"

train_data_pts = get_data_pts(train_path)
test_data_pts = get_data_pts(test_path)

y = get_y(input_shape, y_train_path, train_data_pts, sim=True).astype(np.float32)
y_test = get_y(input_shape, y_test_path, test_data_pts, sim=True).astype(np.float32)

for i in range(repeats):
    for beta in betas:
        fft_path = datasets + f"/translated{choice}{beta}/"
        fft_train_path = fft_path + "train/data"
        fft_test_path = fft_path + "test/data"
        train_data_pts = get_data_pts(fft_train_path)
        test_data_pts = get_data_pts(fft_test_path)

        X = get_X(input_shape, fft_train_path, train_data_pts)
        X_test = get_X(input_shape, fft_test_path, test_data_pts)

        X_unlabelled = load_unlabelled(input_shape, unlabelled_train_path)
        for positive_weight in [1]:
            for discriminator_weight in discr_weights:
                for discriminator_depth in discr_depths:
                    keras.backend.clear_session()
                    gc.collect()
                    model_wrapper = model_functions.unet(img_size=input_shape, num_classes=num_classes)

                    segmentor_training.train_adversarial(model_wrapper,
                                                         X, y, X_unlabelled, X_test, y_test,
                                                         epochs=seg_epochs,
                                                         discriminator_weight=discriminator_weight,
                                                         discriminator_depth=discriminator_depth,
                                                         positive_weight=positive_weight,
                                                         sub_path=f"fft{choice}{beta}discr{sub_path(discriminator_weight, i)}d{discriminator_depth}",
                                                         load_if_possible=False, add_info="",
                                                         done_instructions="",
                                                         save=True)