tf使用GPU训练与推理记录

最新推荐文章于 2024-08-03 10:30:00 发布

buptwhq

最新推荐文章于 2024-08-03 10:30:00 发布

阅读量964

点赞数

分类专栏： tf、pytorch、sklearn Linux Python 文章标签：深度学习机器学习 tensorflow

本文链接：https://blog.csdn.net/whq___/article/details/123532709

版权

Linux 同时被 3 个专栏收录

17 篇文章 7 订阅

订阅专栏

Python

11 篇文章 0 订阅

订阅专栏

tf、pytorch、sklearn

7 篇文章 0 订阅

订阅专栏

本文介绍了如何利用CUDA环境选择GPU设备进行预测，并探讨了在TensorFlow中设置GPU内存增长，实现分布式训练的流程。通过MirroredStrategy进行模型并行训练，提升了模型的训练效率，最后完成模型的保存和资源清理。

摘要由CSDN通过智能技术生成

推理

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1' # 选择GPU设备
predict = model.predict(SignalSeq_16[:10])

训练

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)
mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
    ...
    model = 
    model.compile()
    model.fit()
    model.save()
    del model
    gc.collect()