tensorflow显存自适应分配(tf.1x和tf2.x)

用tensorflow的同学肯定都知道,只要tf的程序运行起来就会默认占满显存,但实际上有些代码使用的数据量远远不需要使用全部显存,这样就造成了浪费,我们可以在代码前加上一些控制语句,使程序占用的显存智能一些。

# 指定显卡
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
......
# tf1.x version  
# # 按比例配置显存
config = tf.ConfigProto()

config.gpu_options.per_process_gpu_memory_fraction = 0.7 # 根据自己的需求确定
# 上面一行还可以改为下面两行
config = tf.ConfigProto(allow_soft_placement=True)
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)  #o.7可以自己修改其他值

with tf.Session(config=config) as sess:

# tf1.x version  
# 自适应显存占用
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:



# tf2.x version  自适应显存占用
# gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
# for gpu in gpus:
#     tf.config.experimental.set_memory_growth(gpu, True)

### TensorFlow 2.7.0 中 `ImageDataGenerator` 的 GPU 配置与使用 #### 正确导入 `ImageDataGenerator` 在 TensorFlow 2.7.0 中,可以通过以下方式导入 `ImageDataGenerator`: ```python from tensorflow.keras.preprocessing.image import ImageDataGenerator ``` 此模块主要用于图像增强预处理操作,适用于小型数据集或需要实时数据增强的情况[^1]。 --- #### 设置 GPU 支持 为了确保 TensorFlow 能够利用 GPU 加速计算,在运行代码之前需确认 GPU 已被正确识别。以下是启用 GPU 的基本步骤: 1. **验证 GPU 可用性** 运行以下代码以检查是否有可用的 GPU 设备: ```python import tensorflow as tf gpus = tf.config.list_physical_devices('GPU') if gpus: print(f"Detected {len(gpus)} GPUs") else: print("No GPU detected.") ``` 2. **限制显存分配** 默认情况下,TensorFlow 会占用所有可用显存。为了避免内存不足问题,可动态调整显存分配策略: ```python for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) ``` 通过上述设置,可以有效减少因显存管理不当而导致的错误[^3]。 --- #### 使用 `ImageDataGenerator` 的示例 下面是一个完整的示例,展示了如何结合 GPU `ImageDataGenerator` 来加载预处理图像数据: ```python import tensorflow as tf from tensorflow.keras.preprocessing.image import ImageDataGenerator # 检查 GPU 是否可用 gpus = tf.config.list_physical_devices('GPU') if gpus: try: # 动态增长显存 for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) except RuntimeError as e: print(e) # 定义 ImageDataGenerator 实例 datagen = ImageDataGenerator( rescale=1./255, # 将像素值缩放到 [0, 1] rotation_range=40, # 随机旋转角度范围 (degrees) width_shift_range=0.2, # 水平随机移动比例 height_shift_range=0.2, # 垂直随机移动比例 shear_range=0.2, # 错切变换强度 zoom_range=0.2, # 缩放比例 horizontal_flip=True # 水平翻转 ) # 从指定目录加载训练数据 train_generator = datagen.flow_from_directory( 'data/train', # 训练数据路径 target_size=(150, 150), # 调整输入图像大小 batch_size=32, class_mode='binary' # 输出标签模式(此处为二分类) ) # 构建简单的 CNN 模型 model = tf.keras.models.Sequential([ tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)), tf.keras.layers.MaxPooling2D((2, 2)), tf.keras.layers.Flatten(), tf.keras.layers.Dense(64, activation='relu'), tf.keras.layers.Dense(1, activation='sigmoid') # 二分类输出层 ]) # 编译模型 model.compile(optimizer=tf.keras.optimizers.Adam(), loss='binary_crossentropy', metrics=['accuracy']) # 开始训练 history = model.fit(train_generator, epochs=10) ``` 在此示例中,`ImageDataGenerator` 提供了多种数据增强功能,而 GPU 则加速了卷积运算其他矩阵乘法操作。 --- #### 注意事项 1. **性能瓶颈分析** 如果发现 GPU 占用率较低,可能是因为 CPU 数据准备速度跟不上 GPU 的计算能力。此时建议改用 `tf.data.Dataset` 接口来替代 `ImageDataGenerator`,因为后者更适合大规模分布式场景下的高效数据流水线[^2]。 2. **版本兼容性** 确保所使用的 CUDA cuDNN 版本与 TensorFlow 2.7.0 兼容。具体依赖关系可通过官方文档查询。 3. **资源监控工具** 在调试过程中,推荐使用 NVIDIA 提供的 `nvidia-smi` 或 TensorBoard 插件监视 GPU 使用情况,以便及时发现问题并优化程序性能[^4]。 ---
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值