forked from Kaggle/docker-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_tensorflow.py
More file actions
74 lines (56 loc) · 2.13 KB
/
test_tensorflow.py
File metadata and controls
74 lines (56 loc) · 2.13 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import unittest
import numpy as np
import tensorflow as tf
from tensorflow.contrib import cudnn_rnn
from tensorflow.python.framework import dtypes
from tensorflow.python.ops import variables
from common import gpu_test
class TestTensorflow(unittest.TestCase):
def test_addition(self):
op = tf.add(2, 3)
sess = tf.Session()
result = sess.run(op)
self.assertEqual(5, result)
def test_conv2d(self):
input = tf.random_normal([1,2,2,1])
filter = tf.random_normal([1,1,1,1])
op = tf.nn.conv2d(input, filter, strides=[1, 1, 1, 1], padding='SAME')
with tf.Session() as sess:
result = sess.run(op)
self.assertEqual(4, len(result.shape))
@gpu_test
def test_cudnn_lstm(self):
num_layers = 4
num_units = 2
batch_size = 8
dir_count = 1
inputs = tf.random_uniform([num_layers * dir_count, batch_size, num_units], dtype=dtypes.float32)
lstm = cudnn_rnn.CudnnLSTM(
num_layers=num_layers,
num_units=num_units,
direction='unidirectional',
kernel_initializer=tf.constant_initializer(0.),
bias_initializer=tf.constant_initializer(0.),
name='test_gru'
)
outputs, _ = lstm(inputs)
total_sum = tf.reduce_sum(outputs)
with tf.Session() as sess:
sess.run(variables.global_variables_initializer())
result = sess.run(total_sum)
self.assertEqual(0, result)
@gpu_test
def test_gpu(self):
with tf.device('/gpu:0'):
m1 = tf.constant([2.0, 3.0], shape=[1, 2], name='a')
m2 = tf.constant([3.0, 4.0], shape=[2, 1], name='b')
op = tf.matmul(m1, m2)
with tf.Session() as sess:
result = sess.run(op)
self.assertEqual(np.array(18, dtype=np.float32, ndmin=2), result)
@gpu_test
def test_is_built_with_cuda(self):
self.assertTrue(tf.test.is_built_with_cuda())
@gpu_test
def test_is_gpu_available(self):
self.assertTrue(tf.test.is_gpu_available(cuda_only=True))