Question

我正在尝试通过此链接重现用于神经样式传递的tensorflow教程： https://www.tensorflow.org/tutorials/generative/style_transfer

我需要在不同的图像上重复执行多次，以创建一个人工数据集。我将转换代码放入一个循环中，但是经过一些迭代后，由于GPU使用量超过15.40GB，我遇到了资源耗尽错误。修改后的代码是这样的：

path_to_photos = '/content/monet_images/JPEGImages/'
style_path = '/content/style.jpg'
total_variation_weight=100

for img in os.listdir(path_to_photos):
  if img[:-3]+'monet.jpg' in os.listdir('/content/drive/My Drive/VOC2007/MonetJPEGImages/'):
    continue
  style_image = load_img(style_path)
  vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
  opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
  style_weight=1e-2
  content_weight=1e4
  # Content layer where will pull our feature maps
  content_layers = ['block5_conv2'] 
  # Style layer of interest
  style_layers = ['block1_conv1',
                  'block2_conv1',
                  'block3_conv1', 
                  'block4_conv1', 
                  'block5_conv1']

  @tf.function()
  def train_step(image):
    with tf.GradientTape() as tape:
      outputs = extractor(image)
      loss = style_content_loss(outputs)
      loss += total_variation_weight*tf.image.total_variation(image)

    grad = tape.gradient(loss, image)
    opt.apply_gradients([(grad, image)])
    image.assign(clip_0_1(image))

  num_content_layers = len(content_layers)
  num_style_layers = len(style_layers)

  style_extractor = vgg_layers(style_layers)
  style_outputs = style_extractor(style_image*255)

  print(img)
  content_path = path_to_photos+img
  content_image = load_img(content_path)
  extractor = StyleContentModel(style_layers, content_layers)
  results = extractor(tf.constant(content_image))
  style_results = results['style']
  style_targets = extractor(style_image)['style']
  content_targets = extractor(content_image)['content']
  image = tf.Variable(content_image)

  start = time.time()

  epochs = 5
  steps_per_epoch = 100

  step = 0
  for n in range(epochs):
    for m in range(steps_per_epoch):
      step += 1
      train_step(image)
    display.clear_output(wait=True)
    print("Train step: {}".format(step))
  end = time.time()
  file_name = img[:-3]+'monet.jpg'
  tensor_to_image(image).save('/content/drive/My Drive/VOC2007/MonetJPEGImages/'+file_name)
  print("Total time: {:.1f}".format(end-start))
  tf.keras.backend.clear_session()

是否有一种方法可以在每次迭代后释放gpu使用情况？预先感谢

Google Colab中Tensorflow的内存不足问题

0 个答案: