""" Leras. like lighter keras. This is my lightweight neural network library written from scratch based on pure tensorflow without keras. Provides: + full freedom of tensorflow operations without keras model's restrictions + easy model operations like in PyTorch, but in graph mode (no eager execution) + convenient and understandable logic Reasons why we cannot import tensorflow or any tensorflow.sub modules right here: 1) program is changing env variables based on DeviceConfig before import tensorflow 2) multiprocesses will import tensorflow every spawn NCHW speed up training for 10-20%. """ import os import sys import warnings warnings.simplefilter(action='ignore', category=FutureWarning) from pathlib import Path import numpy as np # from face_feature.core.interact import interact as io from .device import Devices class nn(): current_DeviceConfig = None tf = None tf_sess = None tf_sess_config = None tf_default_device_name = None data_format = None conv2d_ch_axis = None conv2d_spatial_axes = None floatx = None @staticmethod def initialize(device_config=None, floatx="float32", data_format="NHWC"): if nn.tf is None: if device_config is None: device_config = nn.getCurrentDeviceConfig() nn.setCurrentDeviceConfig(device_config) # Manipulate environment variables before import tensorflow first_run = False # if len(device_config.devices) != 0: # if sys.platform[0:3] == 'win': # # Windows specific env vars # if all( [ x.name == device_config.devices[0].name for x in device_config.devices ] ): # devices_str = "_" + device_config.devices[0].name.replace(' ','_') # else: # devices_str = "" # for device in device_config.devices: # devices_str += "_" + device.name.replace(' ','_') # # compute_cache_path = Path(os.environ['APPDATA']) / 'NVIDIA' / ('ComputeCache' + devices_str) # if not compute_cache_path.exists(): # first_run = True # compute_cache_path.mkdir(parents=True, exist_ok=True) # os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path) # if first_run: # io.log_info("Caching GPU kernels...") import tensorflow tf_version = tensorflow.version.VERSION #if tf_version is None: # tf_version = tensorflow.version.GIT_VERSION if tf_version[0] == 'v': tf_version = tf_version[1:] if tf_version[0] == '2': tf = tensorflow.compat.v1 else: tf = tensorflow import logging # Disable tensorflow warnings tf_logger = logging.getLogger('tensorflow') tf_logger.setLevel(logging.ERROR) if tf_version[0] == '2': tf.disable_v2_behavior() nn.tf = tf # Initialize framework # import core.leras.ops # import core.leras.layers # import core.leras.initializers # import core.leras.optimizers # import core.leras.models # import core.leras.archis # Configure tensorflow session-config if len(device_config.devices) == 0: config = tf.ConfigProto(device_count={'GPU': 0}) nn.tf_default_device_name = '/CPU:0' else: nn.tf_default_device_name = f'/{device_config.devices[0].tf_dev_type}:0' config = tf.ConfigProto() config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices]) config.gpu_options.force_gpu_compatible = True config.gpu_options.allow_growth = True nn.tf_sess_config = config if nn.tf_sess is None: nn.tf_sess = tf.Session(config=nn.tf_sess_config) if floatx == "float32": floatx = nn.tf.float32 elif floatx == "float16": floatx = nn.tf.float16 else: raise ValueError(f"unsupported floatx {floatx}") nn.set_floatx(floatx) nn.set_data_format(data_format) @staticmethod def initialize_main_env(): Devices.initialize_main_env() @staticmethod def set_floatx(tf_dtype): """ set default float type for all layers when dtype is None for them """ nn.floatx = tf_dtype @staticmethod def set_data_format(data_format): if data_format != "NHWC" and data_format != "NCHW": raise ValueError(f"unsupported data_format {data_format}") nn.data_format = data_format if data_format == "NHWC": nn.conv2d_ch_axis = 3 nn.conv2d_spatial_axes = [1,2] elif data_format == "NCHW": nn.conv2d_ch_axis = 1 nn.conv2d_spatial_axes = [2,3] @staticmethod def get4Dshape ( w, h, c ): """ returns 4D shape based on current data_format """ if nn.data_format == "NHWC": return (None,h,w,c) else: return (None,c,h,w) @staticmethod def to_data_format( x, to_data_format, from_data_format): if to_data_format == from_data_format: return x if to_data_format == "NHWC": return np.transpose(x, (0,2,3,1) ) elif to_data_format == "NCHW": return np.transpose(x, (0,3,1,2) ) else: raise ValueError(f"unsupported to_data_format {to_data_format}") @staticmethod def getCurrentDeviceConfig(): if nn.current_DeviceConfig is None: nn.current_DeviceConfig = nn.DeviceConfig.BestGPU() return nn.current_DeviceConfig @staticmethod def setCurrentDeviceConfig(device_config): nn.current_DeviceConfig = device_config @staticmethod def reset_session(): if nn.tf is not None: if nn.tf_sess is not None: nn.tf.reset_default_graph() nn.tf_sess.close() nn.tf_sess = nn.tf.Session(config=nn.tf_sess_config) @staticmethod def close_session(): if nn.tf_sess is not None: nn.tf.reset_default_graph() nn.tf_sess.close() nn.tf_sess = None @staticmethod # def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False): # devices = Devices.getDevices() # if len(devices) == 0: # return [] # # all_devices_indexes = [device.index for device in devices] # # if choose_only_one: # suggest_best_multi_gpu = False # suggest_all_gpu = False # # if suggest_all_gpu: # best_device_indexes = all_devices_indexes # elif suggest_best_multi_gpu: # best_device_indexes = [device.index for device in devices.get_equal_devices(devices.get_best_device()) ] # else: # best_device_indexes = [ devices.get_best_device().index ] # best_device_indexes = ",".join([str(x) for x in best_device_indexes]) # # io.log_info ("") # if choose_only_one: # io.log_info ("Choose one GPU idx.") # else: # io.log_info ("Choose one or several GPU idxs (separated by comma).") # io.log_info ("") # # if allow_cpu: # io.log_info ("[CPU] : CPU") # for device in devices: # io.log_info (f" [{device.index}] : {device.name}") # # io.log_info ("") # # while True: # try: # if choose_only_one: # choosed_idxs = io.input_str("Which GPU index to choose?", best_device_indexes) # else: # choosed_idxs = io.input_str("Which GPU indexes to choose?", best_device_indexes) # # if allow_cpu and choosed_idxs.lower() == "cpu": # choosed_idxs = [] # break # # choosed_idxs = [ int(x) for x in choosed_idxs.split(',') ] # # if choose_only_one: # if len(choosed_idxs) == 1: # break # else: # if all( [idx in all_devices_indexes for idx in choosed_idxs] ): # break # except: # pass # io.log_info ("") # # return choosed_idxs class DeviceConfig(): @staticmethod def ask_choose_device(*args, **kwargs): return nn.DeviceConfig.GPUIndexes( nn.ask_choose_device_idxs(*args,**kwargs) ) def __init__ (self, devices=None): devices = devices or [] if not isinstance(devices, Devices): devices = Devices(devices) self.devices = devices self.cpu_only = len(devices) == 0 @staticmethod def BestGPU(): devices = Devices.getDevices() if len(devices) == 0: return nn.DeviceConfig.CPU() return nn.DeviceConfig([devices.get_best_device()]) @staticmethod def WorstGPU(): devices = Devices.getDevices() if len(devices) == 0: return nn.DeviceConfig.CPU() return nn.DeviceConfig([devices.get_worst_device()]) @staticmethod def GPUIndexes(indexes): if len(indexes) != 0: devices = Devices.getDevices().get_devices_from_index_list(indexes) else: devices = [] return nn.DeviceConfig(devices) @staticmethod def CPU(): return nn.DeviceConfig([])