algorithm_platform/models_processing/model_tf/settings.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :common.py
# @Time      :2025/3/5 16:06
# @Author    :David
# @Company: shenyang JY

import os
import random
import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K


def set_deterministic(seed=42):
    """设置所有可能的随机种子和环境变量"""
    # 设置Python内置随机种子
    random.seed(seed)

    # 设置NumPy随机种子
    np.random.seed(seed)

    # 设置TensorFlow随机种子
    tf.random.set_seed(seed)

    # 设置Keras后端随机种子
    K.set_floatx('float32')
    K.set_epsilon(1e-7)

    # 设置环境变量（影响CUDA和底层操作）
    os.environ['PYTHONHASHSEED'] = str(seed)
    os.environ['TF_DETERMINISTIC_OPS'] = '1'
    os.environ['TF_CUDNN_DETERMINISTIC'] = '1'

    # 配置GPU选项（如果使用GPU）
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
            tf.config.experimental.set_virtual_device_configuration(
                gpus[0],
                [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]
            )
            tf.config.set_soft_device_placement(True)
        except RuntimeError as e:
            print(e)

    # 设置并行线程数
    tf.config.threading.set_inter_op_parallelism_threads(1)
    tf.config.threading.set_intra_op_parallelism_threads(1)


# 在程序初始化时调用
set_deterministic(42)
if __name__ == "__main__":
    run_code = 0