关于TensorFlow GPU 跑神经网络爆内存的处理办法
本人上一篇博客介绍了vgg16模型的finetuning（网址：点击打开链接），在cpu上跑没问题，但是我在全部复制到gpu（gtx750TI）上跑却发生爆内存的情况，下面介绍一下我的解决办法。
不得不承认爆内存跟我的渣渣gpu有关，为了尽量的减少gpu内存的使用率，在加载vgg16模型的时候并没有直接调用命令，而是把网络的权重跟偏置都设为常值，直接赋值，减少了以后sess.run(tf.global_variables_initializer())这句明令时因为run所有的变量而带来的内存消耗。
关于vgg16模型下载上面的链接中有，自行查看。而查看vgg16模型里面的具体代码如下，我们可以看见每一部分的数值：
weight_file = 'vgg/vgg16_weights.npz' #自己改文件的地址
weights = np.load(weight_file)
keys = sorted(weights.keys())
print(weights[keys[29]])#此处查看的是第30个部分的数值 其次关于网上说的batch_size太大也有关系，我电脑这次设置的batch_size为16，实验成功。
整个vgg16在gpu上训练的代码如下：
import tensorflow as tf
import numpy as np
import os
N_CLASSES = 2 # 2个输出神经元，［1，0］ 或者 ［0，1］猫和狗的概率
IMG_W = 224 # 重新定义图片的大小，图片如果过大则训练比较慢
IMG_H = 224
BATCH_SIZE = 16 #每批数据的大小
CAPACITY = 256 #队列容量,队列中最多容纳图片的个数
MAX_STEP = 15000 # 训练的步数，应当 >= 10000
learning_rate = 0.0001 # 学习率，建议刚开始的 learning_rate <= 0.0001
file_dir = 'catdog/train/'
weight_file = 'vgg/vgg16_weights.npz'
weights = np.load(weight_file)
keys = sorted(weights.keys())
#print(weights[keys[0]])
def network(imgs, batch_size, n_classes):
# zero-mean input
with tf.name_scope('preprocess') as scope:
mean = tf.constant([123.68, 116.779, 103.939], dtype=tf.float32, shape=[1, 1, 1, 3], name='img_mean')
images = imgs-mean
# conv1_1
with tf.name_scope('conv1_1') as scope:
kernel = weights[keys[0]]
conv = tf.nn.conv2d(images, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[1]]
out = tf.nn.bias_add(conv, biases)
conv1_1 = tf.nn.relu(out, name=scope)
# conv1_2
with tf.name_scope('conv1_2') as scope:
kernel = weights[keys[2]]
conv = tf.nn.conv2d(conv1_1, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[3]]
out = tf.nn.bias_add(conv, biases)
conv1_2 = tf.nn.relu(out, name=scope)
# pool1
pool1 = tf.nn.max_pool(conv1_2,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool1')
# conv2_1
with tf.name_scope('conv2_1') as scope:
kernel = weights[keys[4]]
conv = tf.nn.conv2d(pool1, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[5]]
out = tf.nn.bias_add(conv, biases)
conv2_1 = tf.nn.relu(out, name=scope)
# conv2_2
with tf.name_scope('conv2_2') as scope:
kernel = weights[keys[6]]
conv = tf.nn.conv2d(conv2_1, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[7]]
out = tf.nn.bias_add(conv, biases)
conv2_2 = tf.nn.relu(out, name=scope)
# pool2
pool2 = tf.nn.max_pool(conv2_2,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool2')
# conv3_1
with tf.name_scope('conv3_1') as scope:
kernel = weights[keys[8]]
conv = tf.nn.conv2d(pool2, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[9]]
out = tf.nn.bias_add(conv, biases)
conv3_1 = tf.nn.relu(out, name=scope)
# conv3_2
with tf.name_scope('conv3_2') as scope:
kernel = weights[keys[10]]
conv = tf.nn.conv2d(conv3_1, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[11]]
out = tf.nn.bias_add(conv, biases)
conv3_2 = tf.nn.relu(out, name=scope)
# conv3_3
with tf.name_scope('conv3_3') as scope:
kernel = weights[keys[12]]
conv = tf.nn.conv2d(conv3_2, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[13]]
out = tf.nn.bias_add(conv, biases)
conv3_3 = tf.nn.relu(out, name=scope)
# pool3
pool3 = tf.nn.max_pool(conv3_3,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool3')
# conv4_1
with tf.name_scope('conv4_1') as scope:
kernel = weights[keys[14]]
conv = tf.nn.conv2d(pool3, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[15]]
out = tf.nn.bias_add(conv, biases)
conv4_1 = tf.nn.relu(out, name=scope)
# conv4_2
with tf.name_scope('conv4_2') as scope:
kernel = weights[keys[16]]
conv = tf.nn.conv2d(conv4_1, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[17]]
out = tf.nn.bias_add(conv, biases)
conv4_2 = tf.nn.relu(out, name=scope)
# conv4_3
with tf.name_scope('conv4_3') as scope:
kernel = weights[keys[18]]
conv = tf.nn.conv2d(conv4_2, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[19]]
out = tf.nn.bias_add(conv, biases)
conv4_3 = tf.nn.relu(out, name=scope)
# pool4
pool4 = tf.nn.max_pool(conv4_3,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool4')
# conv5_1
with tf.name_scope('conv5_1') as scope:
kernel = weights[keys[20]]
conv = tf.nn.conv2d(pool4, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[21]]
out = tf.nn.bias_add(conv, biases)
conv5_1 = tf.nn.relu(out, name=scope)
# conv5_2
with tf.name_scope('conv5_2') as scope:
kernel = weights[keys[22]]
conv = tf.nn.conv2d(conv5_1, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[23]]
out = tf.nn.bias_add(conv, biases)
conv5_2 = tf.nn.relu(out, name=scope)
# conv5_3
with tf.name_scope('conv5_3') as scope:
kernel = weights[keys[24]]
conv = tf.nn.conv2d(conv5_2, kernel, [1, 1, 1, 1], padding='SAME')
biases = weights[keys[25]]
out = tf.nn.bias_add(conv, biases)
conv5_3 = tf.nn.relu(out, name=scope)
# pool5
pool5 = tf.nn.max_pool(conv5_3,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool4')
# fc1
with tf.name_scope('fc1') as scope:
reshape = tf.reshape(pool5, shape=[batch_size, -1])
fc1w = weights[keys[26]]
fc1b = weights[keys[27]]
#pool5_flat = tf.reshape(pool5, [-1, shape])
fc1l = tf.nn.bias_add(tf.matmul(reshape, fc1w), fc1b)
fc1 = tf.nn.relu(fc1l)
# fc2
with tf.name_scope('fc2') as scope:
fc2w = weights[keys[28]]
fc2b = weights[keys[29]]
fc2l = tf.nn.bias_add(tf.matmul(fc1, fc2w), fc2b)
fc2 = tf.nn.relu(fc2l)
# fc3
with tf.name_scope('fc3') as scope:
fc3w = tf.Variable(tf.truncated_normal([4096, n_classes],
dtype=tf.float32,
stddev=1e-1), trainable=True,name='weights')
fc3b = tf.Variable(tf.constant(1.0, shape=[n_classes], dtype=tf.float32),
trainable=True, name='biases')
fc3l = tf.nn.bias_add(tf.matmul(fc2, fc3w), fc3b)
fc3l = tf.nn.softmax(fc3l)
print(fc3l)
return fc3l
# def load_weights(self, weight_file, sess):
# weights = np.load(weight_file)
# keys = sorted(weights.keys())
# for i, k in enumerate(keys):
# if i not in [30,31]:
# print (i, k, np.shape(weights[k]))
# sess.run(self.parameters[i].assign(weights[k]))
# print("-------------all done-------------")
def get_files(file_dir):
# file_dir: 文件夹路径
# return: 乱序后的图片和标签
cats = []
label_cats = []
dogs = []
label_dogs = []
# 载入数据路径并写入标签值
for file in os.listdir(file_dir):
name = file.split('.')
if name[0] == 'cat':
cats.append(file_dir + file) #修改原来的列表,增加对象
label_cats.append(0)
else:
dogs.append(file_dir + file)
label_dogs.append(1)
print("There are %d cats\nThere are %d dogs" % (len(cats), len(dogs)))
# 打乱文件顺序
image_list = np.hstack((cats, dogs)) #把猫放在前面，狗在后面组成一个文件
label_list = np.hstack((label_cats, label_dogs))
temp = np.array([image_list, label_list]) #生成序列
temp = temp.transpose() # 转置
np.random.shuffle(temp)
image_list = list(temp[:, 0]) #下标索引，访问列表值
label_list = list(temp[:, 1])
label_list = [int(i) for i in label_list]
return image_list, label_list
def get_batch(image,label,image_W,image_H,batch_size,capacity):
# 转换数据为 ts 能识别的格式
image = tf.cast(image,tf.string) #将image格式转化为字符串
label = tf.cast(label, tf.int32)
# 将image 和 label 放倒队列里
input_queue = tf.train.slice_input_producer([image,label])
label = input_queue[1]
# 读取图片的全部信息
image_contents = tf.read_file(input_queue[0])
# 把图片解码，channels ＝3 为彩色图片, r，g ，b 黑白图片为 1 ，也可以理解为图片的厚度
image = tf.image.decode_jpeg(image_contents,channels =3)
# 将图片以图片中心进行裁剪或者扩充为 指定的image_W，image_H
image = tf.image.resize_image_with_crop_or_pad(image, image_W, image_H)
# 对数据进行标准化,标准化，就是减去它的均值，除以他的方差
image = tf.image.per_image_standardization(image)
# 生成批次 num_threads 有多少个线程根据电脑配置设置 capacity 队列中 最多容纳图片的个数 tf.train.shuffle_batch 打乱顺序，
image_batch, label_batch = tf.train.batch([image, label],batch_size = batch_size, num_threads = 64, capacity = capacity)
# 重新定义下 label_batch 的形状
label_batch = tf.reshape(label_batch , [batch_size])
# 转化图片格式
image_batch = tf.cast(image_batch,tf.float32)
return image_batch, label_batch
def onehot(labels):
'''''one-hot'''
n_sample = len(labels)
n_class = max(labels)+1
onehot_labels = np.zeros((n_sample ,n_class))
onehot_labels[np.arange(n_sample),labels] = 1
return onehot_labels
def losses(logits, labels):
with tf.variable_scope('loss') as scope:
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits \
(logits=logits, labels=labels, name='xentropy_per_example')
loss = tf.reduce_mean(cross_entropy, name='loss')
return loss
def trainning(loss, learning_rate):
with tf.name_scope('optimizer'):
optimizer = tf.train.AdamOptimizer(learning_rate= learning_rate) #自适应优化器
global_step = tf.Variable(0, name='global_step', trainable=False) #记录全局训练步骤的单值
train_op = optimizer.minimize(loss, global_step= global_step)
return train_op
def evaluation(logits, labels):
with tf.variable_scope('accuracy') as scope:
# k：每个样本的预测结果的前k个最大的数里面是否包含targets预测中的标签，一般都是取1，即取预测最大概率的索引与标签对比。
correct = tf.nn.in_top_k(logits, labels, 1)
correct = tf.cast(correct, tf.float16)
accuracy = tf.reduce_mean(correct)
return accuracy
def run_training():
# 获取图片和标签集
train, train_label = get_files(file_dir)
# 生成批次
train_batch, train_label_batch = get_batch(train,
train_label,
IMG_W,
IMG_H,
BATCH_SIZE,
CAPACITY)
# 进入模型
train_logits = network(train_batch, BATCH_SIZE, N_CLASSES)
# 获取 loss
train_loss = losses(train_logits, train_label_batch)
# 训练
train_op = trainning(train_loss, learning_rate)
# 获取准确率
train__acc = evaluation(train_logits, train_label_batch)
# 合并 summary
sess = tf.Session()
sess.run(tf.global_variables_initializer())
###TensorFlow提供了两个类来帮助多线程的实现：tf.Coordinator和 tf.QueueRunner。
#从设计上这两个类必须被一起使用。Coordinator类可以用来同时停止多个工作线程并且向那个在等待所有工作线程终止的程序报告异常。
#QueueRunner类用来协调多个工作线程同时将多个张量推入同一个队列中
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
try:
for step in np.arange(MAX_STEP):
if coord.should_stop():
break
_, tra_loss, tra_acc = sess.run([train_op, train_loss, train__acc])
if step % 50== 0:
print('Step %d, train loss = %.2f, train accuracy = %.2f%%' %(step, tra_loss, tra_acc*100.0))
except tf.errors.OutOfRangeError:
print('Done training -- epoch limit reached')
finally:
coord.request_stop()
coord.join(threads)
sess.close()
# train
run_training()