def merge_shards(filename, num_shards_to_merge, out_tmp_dir, batch_size):
#print("Merging shards for output file: %s (%d shards to merge)" % (filename,num_shards_to_merge))
tfoptions = TFRecordOptions(TFRecordCompressionType.ZLIB)
record_writer = TFRecordWriter(filename,tfoptions)
binaryInputNCHWPackeds = []
globalInputNCs = []
policyTargetsNCMoves = []
globalTargetsNCs = []
scoreDistrNs = []
selfBonusScoreNs = []
valueTargetsNCHWs = []
for input_idx in range(num_shards_to_merge):
shard_filename = os.path.join(out_tmp_dir, str(input_idx) + ".npz")
#print("Merge loading shard: %d (mem usage %dMB)" % (input_idx,memusage_mb()))
npz = np.load(shard_filename)
assert(set(npz.keys()) == set(keys))
binaryInputNCHWPacked = npz["binaryInputNCHWPacked"]
globalInputNC = npz["globalInputNC"]
policyTargetsNCMove = npz["policyTargetsNCMove"].astype(np.float32)
globalTargetsNC = npz["globalTargetsNC"]
scoreDistrN = npz["scoreDistrN"].astype(np.float32)
selfBonusScoreN = npz["selfBonusScoreN"].astype(np.float32)
valueTargetsNCHW = npz["valueTargetsNCHW"].astype(np.float32)
binaryInputNCHWPackeds.append(binaryInputNCHWPacked)
globalInputNCs.append(globalInputNC)
policyTargetsNCMoves.append(policyTargetsNCMove)
globalTargetsNCs.append(globalTargetsNC)
scoreDistrNs.append(scoreDistrN)
selfBonusScoreNs.append(selfBonusScoreN)
valueTargetsNCHWs.append(valueTargetsNCHW)
###
#WARNING - if adding anything here, also add it to joint_shuffle below!
###
#print("Merge concatenating... (mem usage %dMB)" % memusage_mb())
binaryInputNCHWPacked = np.concatenate(binaryInputNCHWPackeds)
globalInputNC = np.concatenate(globalInputNCs)
policyTargetsNCMove = np.concatenate(policyTargetsNCMoves)
globalTargetsNC = np.concatenate(globalTargetsNCs)
scoreDistrN = np.concatenate(scoreDistrNs)
selfBonusScoreN = np.concatenate(selfBonusScoreNs)
valueTargetsNCHW = np.concatenate(valueTargetsNCHWs)
#print("Merge shuffling... (mem usage %dMB)" % memusage_mb())
joint_shuffle((binaryInputNCHWPacked,globalInputNC,policyTargetsNCMove,globalTargetsNC,scoreDistrN,selfBonusScoreN,valueTargetsNCHW))
#print("Merge writing in batches...")
num_rows = binaryInputNCHWPacked.shape[0]
#Just truncate and lose the batch at the end, it's fine
num_batches = num_rows // batch_size
for i in range(num_batches):
start = i*batch_size
stop = (i+1)*batch_size
example = tfrecordio.make_tf_record_example(
binaryInputNCHWPacked,
globalInputNC,
policyTargetsNCMove,
globalTargetsNC,
scoreDistrN,
selfBonusScoreN,
valueTargetsNCHW,
start,
stop
)
record_writer.write(example.SerializeToString())
jsonfilename = os.path.splitext(filename)[0] + ".json"
with open(jsonfilename,"w") as f:
json.dump({"num_rows":num_rows,"num_batches":num_batches},f)
#print("Merge done %s (%d rows)" % (filename, num_batches * batch_size))
record_writer.close()
return num_batches * batch_size
def merge_shards(filename, num_shards_to_merge, out_tmp_dir, batch_size,
ensure_batch_multiple):
tfoptions = TFRecordOptions(TFRecordCompressionType.ZLIB)
record_writer = TFRecordWriter(filename, tfoptions)
binaryInputNCHWPackeds = []
globalInputNCs = []
policyTargetsNCMoves = []
globalTargetsNCs = []
scoreDistrNs = []
valueTargetsNCHWs = []
for input_idx in range(num_shards_to_merge):
shard_filename = os.path.join(out_tmp_dir, str(input_idx) + ".npz")
with np.load(shard_filename) as npz:
assert (set(npz.keys()) == set(keys))
binaryInputNCHWPacked = npz["binaryInputNCHWPacked"]
globalInputNC = npz["globalInputNC"]
policyTargetsNCMove = npz["policyTargetsNCMove"].astype(np.float32)
globalTargetsNC = npz["globalTargetsNC"]
scoreDistrN = npz["scoreDistrN"].astype(np.float32)
valueTargetsNCHW = npz["valueTargetsNCHW"].astype(np.float32)
binaryInputNCHWPackeds.append(binaryInputNCHWPacked)
globalInputNCs.append(globalInputNC)
policyTargetsNCMoves.append(policyTargetsNCMove)
globalTargetsNCs.append(globalTargetsNC)
scoreDistrNs.append(scoreDistrN)
valueTargetsNCHWs.append(valueTargetsNCHW)
###
#WARNING - if adding anything here, also add it to joint_shuffle below!
###
binaryInputNCHWPacked = np.concatenate(binaryInputNCHWPackeds)
globalInputNC = np.concatenate(globalInputNCs)
policyTargetsNCMove = np.concatenate(policyTargetsNCMoves)
globalTargetsNC = np.concatenate(globalTargetsNCs)
scoreDistrN = np.concatenate(scoreDistrNs)
valueTargetsNCHW = np.concatenate(valueTargetsNCHWs)
joint_shuffle((binaryInputNCHWPacked, globalInputNC, policyTargetsNCMove,
globalTargetsNC, scoreDistrN, valueTargetsNCHW))
num_rows = binaryInputNCHWPacked.shape[0]
#Just truncate and lose the batch at the end, it's fine
num_batches = (
num_rows //
(batch_size * ensure_batch_multiple)) * ensure_batch_multiple
for i in range(num_batches):
start = i * batch_size
stop = (i + 1) * batch_size
example = tfrecordio.make_tf_record_example(
binaryInputNCHWPacked, globalInputNC, policyTargetsNCMove,
globalTargetsNC, scoreDistrN, valueTargetsNCHW, start, stop)
record_writer.write(example.SerializeToString())
jsonfilename = os.path.splitext(filename)[0] + ".json"
with open(jsonfilename, "w") as f:
json.dump({"num_rows": num_rows, "num_batches": num_batches}, f)
record_writer.close()
return num_batches * batch_size
def merge_shards(filename, num_shards_to_merge, out_tmp_dir, batch_size,
ensure_batch_multiple, output_npz):
np.random.seed(
[int.from_bytes(os.urandom(4), byteorder='little') for i in range(5)])
if output_npz:
record_writer = None
else:
tfoptions = TFRecordOptions(TFRecordCompressionType.ZLIB)
record_writer = TFRecordWriter(filename, tfoptions)
binaryInputNCHWPackeds = []
globalInputNCs = []
policyTargetsNCMoves = []
globalTargetsNCs = []
scoreDistrNs = []
valueTargetsNCHWs = []
for input_idx in range(num_shards_to_merge):
shard_filename = os.path.join(out_tmp_dir, str(input_idx) + ".npz")
with np.load(shard_filename) as npz:
assert (set(npz.keys()) == set(keys))
binaryInputNCHWPacked = npz["binaryInputNCHWPacked"]
globalInputNC = npz["globalInputNC"]
policyTargetsNCMove = npz["policyTargetsNCMove"].astype(np.float32)
globalTargetsNC = npz["globalTargetsNC"]
scoreDistrN = npz["scoreDistrN"].astype(np.float32)
valueTargetsNCHW = npz["valueTargetsNCHW"].astype(np.float32)
binaryInputNCHWPackeds.append(binaryInputNCHWPacked)
globalInputNCs.append(globalInputNC)
policyTargetsNCMoves.append(policyTargetsNCMove)
globalTargetsNCs.append(globalTargetsNC)
scoreDistrNs.append(scoreDistrN)
valueTargetsNCHWs.append(valueTargetsNCHW)
###
#WARNING - if adding anything here, also add it to joint_shuffle below!
###
binaryInputNCHWPacked = np.concatenate(binaryInputNCHWPackeds)
globalInputNC = np.concatenate(globalInputNCs)
policyTargetsNCMove = np.concatenate(policyTargetsNCMoves)
globalTargetsNC = np.concatenate(globalTargetsNCs)
scoreDistrN = np.concatenate(scoreDistrNs)
valueTargetsNCHW = np.concatenate(valueTargetsNCHWs)
num_rows = binaryInputNCHWPacked.shape[0]
assert (globalInputNC.shape[0] == num_rows)
assert (policyTargetsNCMove.shape[0] == num_rows)
assert (globalTargetsNC.shape[0] == num_rows)
assert (scoreDistrN.shape[0] == num_rows)
assert (valueTargetsNCHW.shape[0] == num_rows)
[
binaryInputNCHWPacked, globalInputNC, policyTargetsNCMove,
globalTargetsNC, scoreDistrN, valueTargetsNCHW
] = (joint_shuffle_take_first_n(num_rows, [
binaryInputNCHWPacked, globalInputNC, policyTargetsNCMove,
globalTargetsNC, scoreDistrN, valueTargetsNCHW
]))
assert (binaryInputNCHWPacked.shape[0] == num_rows)
assert (globalInputNC.shape[0] == num_rows)
assert (policyTargetsNCMove.shape[0] == num_rows)
assert (globalTargetsNC.shape[0] == num_rows)
assert (scoreDistrN.shape[0] == num_rows)
assert (valueTargetsNCHW.shape[0] == num_rows)
#Just truncate and lose the batch at the end, it's fine
num_batches = (
num_rows //
(batch_size * ensure_batch_multiple)) * ensure_batch_multiple
if output_npz:
start = 0
stop = num_batches * batch_size
np.savez_compressed(
filename,
binaryInputNCHWPacked=binaryInputNCHWPacked[start:stop],
globalInputNC=globalInputNC[start:stop],
policyTargetsNCMove=policyTargetsNCMove[start:stop],
globalTargetsNC=globalTargetsNC[start:stop],
scoreDistrN=scoreDistrN[start:stop],
valueTargetsNCHW=valueTargetsNCHW[start:stop])
else:
for i in range(num_batches):
start = i * batch_size
stop = (i + 1) * batch_size
example = tfrecordio.make_tf_record_example(
binaryInputNCHWPacked, globalInputNC, policyTargetsNCMove,
globalTargetsNC, scoreDistrN, valueTargetsNCHW, start, stop)
record_writer.write(example.SerializeToString())
jsonfilename = os.path.splitext(filename)[0] + ".json"
with open(jsonfilename, "w") as f:
json.dump({"num_rows": num_rows, "num_batches": num_batches}, f)
if record_writer is not None:
record_writer.close()
return num_batches * batch_size