小编cal*_*e90的帖子

在Tensorflow中限制GPU设备

我正在用Python开发一个使用Tensorflow的应用程序和另一个带有GPU的模型。我有一台拥有很多GPU(3xNVIDIA GTX1080)的PC,由于所有模型都尝试使用所有可用的GPU,从而导致OUT_OF_MEMORY_ERROR,我发现可以使用以下命令为Python脚本分配特定的GPU:

os.environ['CUDA_VISIBLE_DEVICES'] = '1'
Run Code Online (Sandbox Code Playgroud)

在这里,我附上我的FCN课程的摘要

class FCN:
  def __init__(self):
    os.environ['CUDA_VISIBLE_DEVICES'] = '1'
    self.keep_probability = tf.placeholder(tf.float32, name="keep_probabilty")
    self.image = tf.placeholder(tf.float32, shape=[None, IMAGE_SIZE, IMAGE_SIZE, 3], name="input_image")
    self.annotation = tf.placeholder(tf.int32, shape=[None, IMAGE_SIZE, IMAGE_SIZE, 1], name="annotation")

    self.pred_annotation, logits = inference(self.image, self.keep_probability)
    tf.summary.image("input_image", self.image, max_outputs=2)
    tf.summary.image("ground_truth", tf.cast(self.annotation, tf.uint8), max_outputs=2)
    tf.summary.image("pred_annotation", tf.cast(self.pred_annotation, tf.uint8), max_outputs=2)
    self.loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
                                                                               labels=tf.squeeze(self.annotation,
                                                                                                 squeeze_dims=[3]),
                                                                               name="entropy")))
    tf.summary.scalar("entropy", self.loss)

...
Run Code Online (Sandbox Code Playgroud)

在同一个文件中FCN.py,我有一个使用该类的main,当Tensorflow打印输出时,我可以看到只使用了1个GPU,正如我期望的那样。

if __name__ == "__main__":
  fcn = FCN()
  fcn.train_model()

  images_dir = '/home/super/datasets/MeterDataset/full-dataset-gas-images/'
  for img_file in os.listdir(images_dir):
    fcn.segment(os.path.join(images_dir, img_file))
Run Code Online (Sandbox Code Playgroud)

输出:

2018-01-09 …
Run Code Online (Sandbox Code Playgroud)

python cuda gpu python-2.7 tensorflow

5
推荐指数
1
解决办法
2529
查看次数

特征提取Tensorflow Estimator

我是使用Tensorflow的新手.

我按照教程介绍了如何为MNIST数据集创建CNN到分类器.我想知道是否有一种方法可以使用Estimator类提取特征,即第一个FC层的值.

这是我的代码.

from __future__ import absolute_import, division, print_function
import tensorflow as tf
import numpy as np

tf.logging.set_verbosity(tf.logging.INFO)

def cnn_model_fn(features, labels, mode):
    # Input Layer
    input_layer = tf.reshape(features["x"], [-1, 28, 28, 1])

    # Convolutional Layer #1
    conv1 = tf.layers.conv2d(inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same", activation=tf.nn.relu)

    # Pooling Layer #1
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)

    # Convolutional Layer #2 and Pooling Layer #2
    conv2 = tf.layers.conv2d(inputs=pool1, filters=64, kernel_size=[5, 5], padding="same", activation=tf.nn.relu)
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)

    # Dense Layer
    pool2_flat …
Run Code Online (Sandbox Code Playgroud)

python tensorflow

2
推荐指数
1
解决办法
268
查看次数

标签 统计

python ×2

tensorflow ×2

cuda ×1

gpu ×1

python-2.7 ×1