我正在用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) 我是使用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)