针对Cloud ML的Google Storage(gs)包装文件输入/输出?

时间:2016-11-03 08:08:56

标签: tensorflow google-cloud-storage google-cloud-ml

谷歌最近宣布了Clould ML,https://cloud.google.com/ml/,这非常有用。但是,一个限制是Tensorflow程序的输入/输出应该支持gs://。

如果我们使用所有tensorflow APIS来读/写文件,它应该没问题,因为这些API支持gs://

但是,如果我们使用open之类的原生文件IO API,则它不起作用,因为他们不理解gs://

例如:

 with open(vocab_file, 'wb') as f:
        cPickle.dump(self.words, f)

此代码无法在Google Cloud ML中使用。

但是,将所有本机文件IO API修改为tensorflow API或Google Storage Python API实际上非常繁琐。有没有简单的方法来做到这一点?任何支持谷歌存储系统的包装器,gs://在本机文件IO之上?

正如Pickled scipy sparse matrix as input data?所建议的那样,也许我们可以使用file_io.read_file_to_string('gs://...'),但这仍然需要重要的代码修改。

3 个答案:

答案 0 :(得分:13)

这样做:

from tensorflow.python.lib.io import file_io

with file_io.FileIO('gs://.....', mode='w+') as f:
    cPickle.dump(self.words, f)

或者您可以像这样阅读pickle文件:

file_stream = file_io.FileIO(train_file, mode='r')
x_train, y_train, x_test, y_test  = pickle.load(file_stream)

答案 1 :(得分:8)

一种解决方案是在程序启动时将所有数据复制到本地磁盘。你可以在运行的Python脚本中使用gsutil来做到这一点,例如:

vocab_file = 'vocab.pickled'
subprocess.check_call(['gsutil', '-m' , 'cp', '-r',
                       os.path.join('gs://path/to/', vocab_file), '/tmp'])

with open(os.path.join('/tmp', vocab_file), 'wb') as f:
  cPickle.dump(self.words, f)

如果您有任何输出,可以将它们写入本地磁盘并gsutil rsync。 (但是,请小心处理正确的重启,因为您可能会被放在不同的机器上。)

另一个解决方案是猴子补丁open(注意:未经测试):

import __builtin__

# NB: not all modes are compatible; should handle more carefully.
# Probably should be reported on
# https://github.com/tensorflow/tensorflow/issues/4357
def new_open(name, mode='r', buffering=-1):
  return file_io.FileIO(name, mode)

__builtin__.open = new_open

请确保在任何模块实际尝试从GCS读取之前执行此操作。

答案 2 :(得分:2)

apache_beam具有gcsio模块,可用于返回标准Python文件对象以读取/写入GCS对象。您可以将此对象与任何适用于Python文件对象的方法一起使用。例如

def open_local_or_gcs(path, mode):
  """Opens the given path."""
  if path.startswith('gs://'):
    try:
      return gcsio.GcsIO().open(path, mode)
    except Exception as e:  # pylint: disable=broad-except
      # Currently we retry exactly once, to work around flaky gcs calls.
      logging.error('Retrying after exception reading gcs file: %s', e)
      time.sleep(10)
      return gcsio.GcsIO().open(path, mode)
  else:
    return open(path, mode)

 with open_local_or_gcs(vocab_file, 'wb') as f:
   cPickle.dump(self.words, f)