text_plain_output_1.png
import numpy as np

import os

import tensorflow as tf

config = {'SEED': 42, 'DEBUG': False, 'test_size': 0.1, 'img_size': 128, 'batch_size': 32, 'num_labels': 0, 'epochs': 10, 'device': 'GPU'}
debug_config = {}

def set_seed(SEED):
    os.environ['PYTHONHASHSEED'] = str(SEED)
    np.random.seed(SEED)
    tf.random.set_seed(SEED)
set_seed(config['SEED'])

def get_device(device):
    if device == 'TPU':
        try:
            tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect()
            strategy = tf.distribute.TPUStrategy(tpu)
        except ValueError:
            print('Cannot initialize TPU')
    if device == 'GPU':
        strategy = tf.distribute.MirroredStrategy()
    print('Number of accelerators: ', strategy.num_replicas_in_sync)
    return strategy
strategy = get_device(config['device'])
config['batch_size'] = config['batch_size'] * strategy.num_replicas_in_sync