split.py

import sys
import os.path
import numpy as np

n = int(sys.argv[1])
assert 1 < n < 100
fn = sys.argv[2]
pattern = fn.rsplit(".", 1)[0] + "_%02d.npy"

x = np.load(fn)

nx = x.shape[0] / n
assert(x.shape[0] % n == 0);
print "examples per shard", nx

for i in range(n):
  sfn = pattern % i
  s = x[nx*i: nx*(i+1)]
  np.save(sfn, s)
  print i, sfn