前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >Tensorflow 测试一段能运行在 GPU 的代码

Tensorflow 测试一段能运行在 GPU 的代码

作者头像
runzhliu
发布2020-08-06 10:43:29
3.5K0
发布2020-08-06 10:43:29
举报
文章被收录于专栏:容器计算

1 Overview

官方文档「又长又臭」,我只是想在 Kubernetes 集群里,运行一个能跑在 GPU 显卡的程序而已,文档太多,看的眼花缭乱,本文就讲一个简单的例子。

2 Example

例子来源于 gihub 上的一段 code,test_single_gpu.py,核心代码很简单,就是在第一块 GPU 上做一个矩阵的运算。

代码语言:javascript
复制
'''
Single GPU computing
'''
with tf.device('/gpu:0'):
    a = tf.placeholder(tf.float32, [10000, 10000])
    b = tf.placeholder(tf.float32, [10000, 10000])
    # Compute A^n and B^n and store results in c1
    c1.append(matpow(a, n))
    c1.append(matpow(b, n))

with tf.device('/cpu:0'):
  sum = tf.add_n(c1) #Addition of all elements in c1, i.e. A^n + B^n

t1_1 = datetime.datetime.now()
with tf.Session(config=tf.ConfigProto(log_device_placement=log_device_placement)) as sess:
    # Run the op.
    sess.run(sum, {a:A, b:B})
t2_1 = datetime.datetime.now()

Config 里的 log_device_placement 是用来将设备上对 Tensor 的各种操作打印出来。

Enabling device placement logging causes any Tensor allocations or operations to be printed.

然后将这份代码放到 Tensorflow 的官方镜像里,docker build 一下,记得要选 GPU 的镜像,否则没有 CUDA 这些库是跑步起来的。

代码语言:javascript
复制
FROM tensorflow/tensorflow:1.14.0-gpu-py3
COPY test_single_gpu.py /
CMD ["python", "/test_single_gpu.py"]

在 Kubernetes 里运行一个。

代码语言:javascript
复制
apiVersion: v1
kind: Pod
metadata:
  name: tensorflow-gpu
  labels:
    app: tensorflow-gpu
spec:
  containers:
  - name: tensorflow-gpu
    image: tensorflow-gpu-test

3 Summary

测试一段 GPU 的代码,将代码放到合适版本的 Tenorflow 官方的 GPU 镜像,然后通过 Kubernetes 运行起来即可,当然其中需要配置好的 nvidia-docker 之类的环境,本文就不多赘述了。

本文参与 腾讯云自媒体同步曝光计划,分享自作者个人站点/博客。
原始发表:2020/04/08 ,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 作者个人站点/博客 前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 1 Overview
  • 2 Example
  • 3 Summary
相关产品与服务
容器服务
腾讯云容器服务(Tencent Kubernetes Engine, TKE)基于原生 kubernetes 提供以容器为核心的、高度可扩展的高性能容器管理服务,覆盖 Serverless、边缘计算、分布式云等多种业务部署场景,业内首创单个集群兼容多种计算节点的容器资源管理模式。同时产品作为云原生 Finops 领先布道者,主导开源项目Crane,全面助力客户实现资源优化、成本控制。
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档