setup_data.py 文件源码

def main():
  data_dir = sys.argv[1]
  model_dir = sys.argv[2]
  rep_dir = sys.argv[3]
  gpu_mem = sys.argv[4]
  main_mem = sys.argv[5]
  numsplits = int(sys.argv[6])

  data_pbtxt_file = os.path.join(data_dir, 'flickr.pbtxt')
  data_pb = util.ReadData(data_pbtxt_file)
  EditPaths(data_pb, data_dir, gpu_mem, main_mem)
  with open(data_pbtxt_file, 'w') as f:
    text_format.PrintMessage(data_pb, f)
  EditTrainers(data_dir, model_dir, rep_dir, numsplits)

  data_pbtxt_file_z = os.path.join(data_dir, 'flickr_z.pbtxt')
  data_pbtxt_file_nnz = os.path.join(data_dir, 'flickr_nnz.pbtxt')
  if not os.path.exists(data_pbtxt_file_z):
    CreateMissingTextData(data_pb, data_pbtxt_file_z, data_pbtxt_file_nnz)
  data_pb = util.ReadData(data_pbtxt_file_z)
  EditPaths(data_pb, data_dir, gpu_mem, main_mem)
  with open(data_pbtxt_file_z, 'w') as f:
    text_format.PrintMessage(data_pb, f)
  data_pb = util.ReadData(data_pbtxt_file_nnz)
  EditPaths(data_pb, data_dir, gpu_mem, main_mem)
  with open(data_pbtxt_file_nnz, 'w') as f:
    text_format.PrintMessage(data_pb, f)