zhengyang-wang / deeplab-v2--resnet-101--tensorflow Goto Github PK

-----------build encoder: deeplab pre-trained-----------
after start block: (1, ?, ?, 64)
after block1: (1, ?, ?, 256)
after block2: (1, ?, ?, 512)
after block3: (1, ?, ?, 1024)
after block4: (1, ?, ?, 2048)
-----------build decoder-----------
after aspp block: (1, ?, ?, 2)

flags.DEFINE_float('momentum', 0.9, 'momentum')
	flags.DEFINE_string('encoder_name', 'res101', 'name of pre-trained model, res101, res50 or deeplab')
	flags.DEFINE_string('pretrain_file', './reference model/resnet_v1_101.ckpt', 'pre-trained model filename corresponding to encoder_name')
	flags.DEFINE_string('data_list', './dataset_cityscapes/train_fine.txt', 'training data list filename')
...
	flags.DEFINE_integer('input_height', 713, 'input image height')
	flags.DEFINE_integer('input_width', 713, 'input image width')
	flags.DEFINE_integer('num_classes', 19, 'number of classes')

step 0 	 loss = 6.337, (12.625 sec/step)
step 1 	 loss = 6.133, (1.820 sec/step)
step 2 	 loss = 2.675, (1.625 sec/step)
step 3 	 loss = 6.042, (1.630 sec/step)
step 4 	 loss = 15.278, (1.545 sec/step)
step 5 	 loss = 12.153, (1.532 sec/step)
step 6 	 loss = 52.724, (1.100 sec/step)
step 7 	 loss = 940.443, (1.041 sec/step)
step 8 	 loss = 5393914151199927058379743980628738048.000, (1.052 sec/step)
step 9 	 loss = nan, (1.093 sec/step)
step 10 	 loss = nan, (1.445 sec/step)
step 11 	 loss = nan, (1.025 sec/step)

	with tf.variable_scope(scope_name) as scope:
		outputs = self._start_block('conv0')
		print("after start block:", outputs.shape)
		with tf.variable_scope('group0') as scope:
			outputs = self._bottleneck_resblock(outputs, 256, 'block0',	identity_connection=False)
			outputs = self._bottleneck_resblock(outputs, 256, 'block1')
			outputs = self._bottleneck_resblock(outputs, 256, 'block2')
			print("after group0 :", outputs.shape)

	o_b2b = self._conv2d(o_b2a, 3, num_o / 4, 1, name='%s/conv2' % name)
	o_b2b = self._batch_norm(o_b2b, name='%s/conv2' % name, is_training=False, activation_fn=tf.nn.relu)

	o_b2c = self._conv2d(o_b2b, 1, num_o, 1, name='%s/conv3' % name)
	o_b2c = self._batch_norm(o_b2c, name='%s/conv3' % name, is_training=False, activation_fn=None)
	# add
	outputs = self._add([o_b1,o_b2c], name='%s/add' % name)
	# relu
	outputs = self._relu(outputs, name='%s/relu' % name)
	return outputs

_R_MEAN = 123.68 / 255
_G_MEAN = 116.78 / 255
_B_MEAN = 103.94 / 255

channels = tf.split(axis=2, num_or_size_splits=num_channels, value=image)
  for i in range(num_channels):
    channels[i] -= means[i]
  return tf.concat(axis=2, values=channels)

 restore_var = [v for v in tf.global_variables() if 'fc' not in v.name]
# Trainable Variables
all_trainable = tf.trainable_variables()
# Fine-tune part
encoder_trainable = [v for v in all_trainable if 'fc' not in v.name] # lr * 1.0
# Decoder part
decoder_trainable = [v for v in all_trainable if 'fc' in v.name]
		....
decoder_w_trainable = [v for v in decoder_trainable if 'weights' in v.name or 'gamma' in v.name] # lr * 10 

decoder_b_trainable = [v for v in decoder_trainable if 'biases' in v.name or 'beta' in v.name] # lr * 20.0

o=self._conv2d_bn(x, 1, 256, 1, name='fc1', biased=True)
o_bn = self._batch_norm(o, name='fc1_bn', is_training=self.phase, trainable=self.phase,activation_fn=tf.nn.relu)

decoder_w_trainable = [v for v in decoder_trainable if 'weights' in v.name or 'gamma' in v.name] # lr * 10.0
decoder_b_trainable = [v for v in decoder_trainable if 'biases' in v.name or 'beta' in v.name] # lr * 20.0
		

"""Run DeepLab-ResNet on a given image.

This script computes a segmentation mask for a given image.
"""

from __future__ import print_function

import argparse
import os
from PIL import Image
from network import *
from utils import decode_labels


IMG_MEAN = np.array((103.939, 116.779, 123.68), dtype=np.float32)

NUM_CLASSES = 19
SAVE_DIR = './output/'


def get_arguments():
    """Parse all the arguments provided from the CLI.

    Returns:
      A list of parsed arguments.
    """
    parser = argparse.ArgumentParser(description="DeepLabLFOV Network Inference.")
    parser.add_argument("--img_path", type=str,default='./input/frankfurt_000000_000294_leftImg8bit.png',
                        help="Path to the RGB image file.")
    parser.add_argument("--model_weights", type=str, default='./model_cityscape/model.ckpt-10000',
                        help="Path to the file with model weights.")
    parser.add_argument("--num-classes", type=int, default=NUM_CLASSES,
                        help="Number of classes to predict (including background).")
    parser.add_argument("--save-dir", type=str, default=SAVE_DIR,
                        help="Where to save predicted mask.")
    return parser.parse_args()


def load(saver, sess, ckpt_path):
    '''Load trained weights.

    Args:
      saver: TensorFlow saver object.
      sess: TensorFlow session.
      ckpt_path: path to checkpoint file with parameters.
    '''
    saver.restore(sess, ckpt_path)
    print("Restored model parameters from {}".format(ckpt_path))


def main():
    """Create the model and start the evaluation process."""
    args = get_arguments()

    # Prepare image.
    img = tf.image.decode_jpeg(tf.read_file(args.img_path), channels=3)
    # Convert RGB to BGR.
    img_r, img_g, img_b = tf.split(axis=2, num_or_size_splits=3, value=img)
    img = tf.cast(tf.concat(axis=2, values=[img_b, img_g, img_r]), dtype=tf.float32)
    # Extract mean.
    img -= IMG_MEAN

    # Create network. Deeplab_v2(self.image_batch, self.conf.num_classes, False)
    net = Deeplab_v2(tf.expand_dims(img, dim=0), args.num_classes, False)

    # Which variables to load.
    restore_var = tf.global_variables()

    # Predictions.
    raw_output = net.outputs
    raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(img)[0:2, ])
    raw_output_up = tf.argmax(raw_output_up, dimension=3)
    pred = tf.expand_dims(raw_output_up, dim=3)

    # Set up TF session and initialize variables.
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()

    sess.run(init)

    # Load weights.
    loader = tf.train.Saver(var_list=restore_var)
    load(loader, sess, args.model_weights)

    # Perform inference.
    preds = sess.run(pred)

    msk = decode_labels(preds, num_classes=args.num_classes)
    im = Image.fromarray(msk[0])
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    im.save(args.save_dir + 'mask.png')

    print('The output file has been saved to {}'.format(args.save_dir + 'mask.png'))


if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    main()

label_colours = [(128, 64, 128), (244, 35, 231), (69, 69, 69)
                # 0 = road, 1 = sidewalk, 2 = building
                ,(102, 102, 156), (190, 153, 153), (153, 153, 153)
                # 3 = wall, 4 = fence, 5 = pole
                ,(250, 170, 29), (219, 219, 0), (106, 142, 35)
                # 6 = traffic light, 7 = traffic sign, 8 = vegetation
                ,(152, 250, 152), (69, 129, 180), (219, 19, 60)
                # 9 = terrain, 10 = sky, 11 = person
                ,(255, 0, 0), (0, 0, 142), (0, 0, 69)
                # 12 = rider, 13 = car, 14 = truck
                ,(0, 60, 100), (0, 79, 100), (0, 0, 230)
                # 15 = bus, 16 = train, 17 = motocycle
                ,(119, 10, 32), (1, 1, 1)]
                # 18 = bicycle, 19 = void label

       # training
	flags.DEFINE_integer('num_steps', 20000, 'maximum number of iterations')
	flags.DEFINE_integer('save_interval', 1000, 'number of iterations for saving and visualization')
	flags.DEFINE_integer('random_seed', 1234, 'random seed')
	flags.DEFINE_float('weight_decay', 0.0005, 'weight decay rate')
	flags.DEFINE_float('learning_rate', 2.5e-4, 'learning rate')
	flags.DEFINE_float('power', 0.9, 'hyperparameter for poly learning rate')
	flags.DEFINE_float('momentum', 0.9, 'momentum')
	flags.DEFINE_string('encoder_name', 'deeplab', 'name of pre-trained model, res101, res50 or deeplab')
	flags.DEFINE_string('pretrain_file', './reference model/deeplab_resnet_init.ckpt', 'pre-trained model filename corresponding to encoder_name')
	flags.DEFINE_string('data_list', './dataset_cityscapes/train_fine.txt', 'training data list filename')

	# testing / validation
	flags.DEFINE_integer('valid_step', 2000, 'checkpoint number for testing/validation')
	flags.DEFINE_integer('valid_num_steps', 1449, '= number of testing/validation samples')
	flags.DEFINE_string('valid_data_list', './dataset_cityscapes/val_fine.txt', 'testing/validation data list filename')

	# data
	flags.DEFINE_string('data_dir', './cityscapes/leftImg8bit_trainvaltest', 'data directory')
	flags.DEFINE_integer('batch_size', 2, 'training batch size')
	flags.DEFINE_integer('input_height', 512, 'input image height')
	flags.DEFINE_integer('input_width', 1024, 'input image width')
	flags.DEFINE_integer('num_classes', 19, 'number of classes')
	flags.DEFINE_integer('ignore_label', 255, 'label pixel value that should be ignored')
	flags.DEFINE_boolean('random_scale', True, 'whether to perform random scaling data-augmentation')
	flags.DEFINE_boolean('random_mirror', True, 'whether to perform random left-right flipping data-augmentation')	
	```

 #mIoU
pred_logits = tf.reshape(self.pred, [-1, ])
gt = tf.reshape(self.label_batch, [-1, ])
 # Ignoring all labels greater than or equal to n_classes.
temp = tf.less_equal(gt, self.conf.num_classes - 1)
weights = tf.cast(temp, tf.int32)
# fix for tf 1.3.0
gt = tf.where(temp, gt, tf.cast(temp, tf.uint8))        
 self.mIoU, self.mIou_update_op = tf.contrib.metrics.streaming_mean_iou(pred_logits, gt, num_classes=self.conf.num_classes, weights=weights)

self.sess.run(tf.local_variables_initializer())

                loss_value, images, labels, preds, summary, _,_, = self.sess.run(
                    [self.reduced_loss,
                    self.image_batch,
                    self.label_batch,
                    self.pred,
                    self.total_summary,
                    self.train_op,
                    self.mIou_update_op],
                    feed_dict=feed_dict)
m_IoU = self.mIoU.eval(session=self.sess)

print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step), Mean IoU: {:.3f}'.format(step, loss_value, duration, m_IoU))

step 0 	 loss = 4.039, (3.927 sec/step), Mean IoU: 0.004
step 1 	 loss = 3.142, (0.893 sec/step), Mean IoU: 0.004
step 2 	 loss = 2.124, (0.660 sec/step), Mean IoU: 0.004
step 3 	 loss = 1.549, (0.655 sec/step), Mean IoU: 0.004
step 4 	 loss = 2.705, (0.677 sec/step), Mean IoU: 0.004
step 5 	 loss = 1.956, (0.700 sec/step), Mean IoU: 0.004
...
step 502 	 loss = 0.769, (0.852 sec/step), Mean IoU: 0.004
step 503 	 loss = 0.525, (0.831 sec/step), Mean IoU: 0.004
step 504 	 loss = 1.202, (0.812 sec/step), Mean IoU: 0.004

sheldon@amax:~/Projects/Deeplab-v2--ResNet-101--Tensorflow$ python3 main.py
2017-12-28 11:41:44.924418: I tensorflow/core/platform/cpu_feature_guard.cc:137] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA
2017-12-28 11:41:45.623542: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1030] Found device 0 with properties:
name: GeForce GTX 1080 Ti major: 6 minor: 1 memoryClockRate(GHz): 1.582
pciBusID: 0000:8a:00.0
totalMemory: 10.91GiB freeMemory: 10.75GiB
2017-12-28 11:41:45.623612: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1120] Creating TensorFlow device (/device:GPU:0) -> (device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:8a:00.0, compute capability: 6.1)
-----------build encoder: deeplab pre-trained-----------
after start block: (10, 81, 81, 64)
after block1: (10, 81, 81, 256)
after block2: (10, 41, 41, 512)
after block3: (10, 41, 41, 1024)
after block4: (10, 41, 41, 2048)
-----------build decoder-----------
after aspp block: (10, 41, 41, 21)
Restored model parameters from /data2/deeplab_resnet_init.ckpt
2017-12-28 11:42:09.929249: W tensorflow/core/common_runtime/bfc_allocator.cc:273] Allocator (GPU_0_bfc) ran out of memory trying to allocate 720.00MiB.  Current allocation summary follows.
2017-12-28 11:42:09.929406: I tensorflow/core/common_runtime/bfc_allocator.cc:627] Bin (256):   Total Chunks: 278, Chunks in use: 216. 69.5KiB allocated for chunks. 54.0KiB in use in bin. 15.6KiB client-requested in use in bin.
2017-12-28 11:42:09.929433: I tensorflow/core/common_runtime/bfc_allocator.cc:627] Bin (512):   Total Chunks: 65, Chunks in use: 64. 32.5KiB allocated for chunks. 32.0KiB in use in bin. 32.0KiB client-requested in use in bin.
2017-12-28 11:42:09.929452: I tensorflow/core/common_runtime/bfc_allocator.cc:627] Bin (1024):  Total Chunks: 401, Chunks in use: 401. 401.2KiB allocated for chunks. 401.2KiB in use in bin. 401.0KiB client-requested in use in bin.
2017-12-28 11:42:09.929471: I tensorflow/core/common_runtime/bfc_allocator.cc:627] Bin (2048):  Total Chunks: 88, Chunks in use: 88. 176.0KiB allocated for chunks. 176.0KiB in use in bin. 176.0KiB client-requested in use in bin.

> (There were way too much similar outputs, so I just left out most of the lines here)

2017-12-28 11:42:09.950169: I tensorflow/core/common_runtime/bfc_allocator.cc:679] 1 Chunks of size 424673280 totalling 405.00MiB
2017-12-28 11:42:09.950182: I tensorflow/core/common_runtime/bfc_allocator.cc:679] 1 Chunks of size 663552000 totalling 632.81MiB
2017-12-28 11:42:09.950194: I tensorflow/core/common_runtime/bfc_allocator.cc:683] Sum Total of in-use chunks: 9.55GiB
2017-12-28 11:42:09.950211: I tensorflow/core/common_runtime/bfc_allocator.cc:685] Stats:
Limit:                 10968825856
InUse:                 10253331200
MaxInUse:              10265177344
NumAllocs:                    4856
MaxAllocSize:            802160640

2017-12-28 11:42:09.950341: W tensorflow/core/common_runtime/bfc_allocator.cc:277] **********************************************************************************************______
2017-12-28 11:42:09.950375: W tensorflow/core/framework/op_kernel.cc:1192] Resource exhausted: OOM when allocating tensor with shape[5760,4,4,2048]
2017-12-28 11:42:09.973697: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 2.00GiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
2017-12-28 11:42:09.973763: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 4.00MiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
2017-12-28 11:42:09.973796: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 928.77MiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
2017-12-28 11:42:09.999611: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.58GiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
2017-12-28 11:42:09.999662: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 4.00MiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
2017-12-28 11:42:09.999689: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 415.06MiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
2017-12-28 11:42:10.017460: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.21GiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
2017-12-28 11:42:10.017501: W tensorflow/core/common_runtime/bfc_allocator.cc:217] Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.67GiB. The caller indicates that this is not a failure, but may mean that there could be performance gains if more memory is available.
Traceback (most recent call last):
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/client/session.py", line 1323, in _do_call
    return fn(*args)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/client/session.py", line 1302, in _run_fn
    status, run_metadata)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/framework/errors_impl.py", line 473, in __exit__
    c_api.TF_GetCode(self.status.status))
tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[5760,4,4,2048]
         [[Node: fc1_voc12_c3/convolution/SpaceToBatchND = SpaceToBatchND[T=DT_FLOAT, Tblock_shape=DT_INT32, Tpaddings=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](res5c_relu, fc1_voc12_c3/convolution/SpaceToBatchND/block_shape, fc1_voc12_c3/convolution/SpaceToBatchND/paddings)]]
         [[Node: add/_1131 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_5776_add", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "main.py", line 82, in <module>
    tf.app.run()
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/platform/app.py", line 48, in run
    _sys.exit(main(_sys.argv[:1] + flags_passthrough))
  File "main.py", line 76, in main
    getattr(model, args.option)()
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/model.py", line 60, in train
    feed_dict=feed_dict)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/client/session.py", line 889, in run
    run_metadata_ptr)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/client/session.py", line 1120, in _run
    feed_dict_tensor, options, run_metadata)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/client/session.py", line 1317, in _do_run
    options, run_metadata)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/client/session.py", line 1336, in _do_call
    raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape[5760,4,4,2048]
         [[Node: fc1_voc12_c3/convolution/SpaceToBatchND = SpaceToBatchND[T=DT_FLOAT, Tblock_shape=DT_INT32, Tpaddings=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](res5c_relu, fc1_voc12_c3/convolution/SpaceToBatchND/block_shape, fc1_voc12_c3/convolution/SpaceToBatchND/paddings)]]
         [[Node: add/_1131 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_5776_add", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

Caused by op 'fc1_voc12_c3/convolution/SpaceToBatchND', defined at:
  File "main.py", line 82, in <module>
    tf.app.run()
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/platform/app.py", line 48, in run
    _sys.exit(main(_sys.argv[:1] + flags_passthrough))
  File "main.py", line 76, in main
    getattr(model, args.option)()
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/model.py", line 36, in train
    self.train_setup()
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/model.py", line 177, in train_setup
    net = Deeplab_v2(self.image_batch, self.conf.num_classes, True)
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/network.py", line 34, in __init__
    self.build_network()
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/network.py", line 38, in build_network
    self.outputs = self.build_decoder(self.encoding)
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/network.py", line 64, in build_decoder
    outputs = self._ASPP(encoding, self.num_classes, [6, 12, 18, 24])
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/network.py", line 125, in _ASPP
    o.append(self._dilated_conv2d(x, 3, num_o, d, name='fc1_voc12_c%d' % i, biased=True))
  File "/home/sheldon/Projects/Deeplab-v2--ResNet-101--Tensorflow/network.py", line 150, in _dilated_conv2d
    o = tf.nn.atrous_conv2d(x, w, dilation_factor, padding='SAME')
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/ops/nn_ops.py", line 1137, in atrous_conv2d
    name=name)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/ops/nn_ops.py", line 751, in convolution
    return op(input, filter)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/ops/nn_ops.py", line 835, in __call__
    return self.conv_op(inp, filter)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/ops/nn_ops.py", line 499, in __call__
    return self.call(inp, filter)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/ops/nn_ops.py", line 490, in _with_space_to_batch_call
    paddings=paddings)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/ops/gen_array_ops.py", line 4922, in space_to_batch_nd
    paddings=paddings, name=name)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
    op_def=op_def)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/framework/ops.py", line 2956, in create_op
    op_def=op_def)
  File "/home/sheldon/.local/lib/python3.5/site-packages/tensorflow/python/framework/ops.py", line 1470, in __init__
    self._traceback = self._graph._extract_stack()  # pylint: disable=protected-access

ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[5760,4,4,2048]
         [[Node: fc1_voc12_c3/convolution/SpaceToBatchND = SpaceToBatchND[T=DT_FLOAT, Tblock_shape=DT_INT32, Tpaddings=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](res5c_relu, fc1_voc12_c3/convolution/SpaceToBatchND/block_shape, fc1_voc12_c3/convolution/SpaceToBatchND/paddings)]]
         [[Node: add/_1131 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_5776_add", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]