def load_node(node_name, tag=None, iter=None):
filename = node_name
if tag is not None:
filename += '_' + str(tag)
if iter is not None:
filename += ('_%05d' % iter)
filepath = os.path.join(Paths.weight_dir_path(''), filename)
gnode = GNode.load(filepath, Model=ImgGAN)
return gnode
def save_node(node, tag=None, iter=None):
# type: (GNode, str, int) -> None
filename = node.name
if tag is not None:
filename += '_' + str(tag)
if iter is not None:
filename += ('_%05d' % iter)
filename = filename + '.pt'
filepath = os.path.join(Paths.weight_dir_path(''), filename)
node.save(filepath)
H.batch_size, H.classes)
x_seed, l_seed = dl.random_batch('test', 512)
tree = GanTree('gtree', ImgGAN, H, x_seed)
root = tree.create_child_node(dist_params, gan)
root.set_trainer(dl, H, train_config, Model=GanImgTrainer)
root.train(H.root_gan_iters)
# root.gan.load_params(None, None, path = '../experiments/GANTree_Fashion_MNIST_Mixed_trial1/weights/iter/node0/iter-80000.pt')
# root = GNode.load('../experiments/GANTree_Fashion_MNIST_Mixed_trial1/weights/best_node_0_mnist_phase_1_2.pickle', root)
dl_set = {0: dl}
leaf_nodes = {0}
bash_utils.create_dir(Paths.weight_dir_path(''), log_flag=False)
for i_modes in range(4):
node_id = find_next_node()
logger.info(
colored('Next Node to split: %d' % node_id, 'green', attrs=['bold']))
node = tree.nodes[node_id]
x_seed, l_seed = dl_set[node_id].random_batch('test', 512)
train_node(node,
phase1_epochs=H.phase1_epochs,
phase2_iters=H.phase2_iters)
