def create_sgd_optimizer(hparams: tf.contrib.training.HParams):
"""Creates an optimizer as specified in hparams."""
num_epochs = graph_builders.get_or_create_num_epochs()
learning_rates = hparams.learning_rates
learning_rate_stops = hparams.learning_rate_stops
learning_rate = tf.train.piecewise_constant(
num_epochs, learning_rate_stops, learning_rates)
return OPTIMIZERS[hparams.optimizer](learning_rate, beta2=hparams.beta2)
def build_opt_ops(
self,
wavefunction: wavefunctions.Wavefunction,
target_wavefunction: wavefunctions.Wavefunction,
hparams: tf.contrib.training.HParams,
shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
) -> NamedTuple:
"""Adds wavefunction optimization ops to the graph.
Args:
wavefunction: Wavefunction model to optimize.
target_wavefunction: Wavefunction model that serves as a target state.
shared_resources: Resources sharable among different modules.
hparams: Hyperparameters of the optimization procedure.
Returns:
NamedTuple holding tensors needed to run a training epoch.
"""
del shared_resources # Not used by BasisIterationSWO.
batch_size = hparams.batch_size
n_sites = hparams.num_sites
basis_dataset = tf.contrib.data.CsvDataset(
hparams.basis_file_path, [tf.float32 for _ in range(n_sites)],
header=False, field_delim=' ')
basis_dataset = basis_dataset.map(lambda *x: tf.convert_to_tensor(x))
shuffle_batch = scipy.special.binomi(n_sites, n_sites / 2)
basis_dataset = basis_dataset.shuffle(shuffle_batch)
basis_dataset = basis_dataset.batch(batch_size)
basis_dataset = basis_dataset.repeat()
config_iterator = basis_dataset.make_one_shot_iterator()
configs = config_iterator.get_next() * 2. - 1.
psi = wavefunction(configs)
psi_target = target_wavefunction(configs)
loss = tf.reduce_mean(
tf.squared_difference(psi, psi_target * np.sqrt(2 ** n_sites)))
opt_v = wavefunction.get_trainable_variables()
optimizer = create_sgd_optimizer(hparams)
train_step = optimizer.minimize(loss, var_list=opt_v)
num_epochs = graph_builders.get_or_create_num_epochs()
epoch_increment = tf.assign_add(num_epochs, 1)
train_ops = TrainOpsSupervised(
accumulate_gradients=None,
apply_gradients=train_step,
reset_gradients=None,
mc_step=None,
acc_rate=None,
metrics=loss,
update_wf_norm=None,
epoch_increment=epoch_increment,
)
return train_ops
def build_opt_ops(
self,
wavefunction: wavefunctions.Wavefunction,
target_wavefunction: wavefunctions.Wavefunction,
hparams: tf.contrib.training.HParams,
shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
) -> NamedTuple:
"""Adds wavefunction optimization ops to the graph.
Args:
wavefunction: Wavefunction model to optimize.
target_wavefunction: Wavefunction model that serves as a target state.
shared_resources: Resources sharable among different modules.
hparams: Hyperparameters of the optimization procedure.
Returns:
NamedTuple holding tensors needed to run a training epoch.
"""
batch_size = hparams.batch_size
n_sites = hparams.num_sites
configs = graph_builders.get_configs(shared_resources, batch_size, n_sites)
mc_step, acc_rate = graph_builders.get_monte_carlo_sampling(
shared_resources, configs, wavefunction)
psi = wavefunction(configs)
psi_target = target_wavefunction(configs)
loss = tf.reduce_mean(
tf.squared_difference(psi, psi_target * np.sqrt(2**n_sites)) /
(tf.square(tf.stop_gradient(psi)))
)
opt_v = wavefunction.get_trainable_variables()
optimizer = create_sgd_optimizer(hparams)
train_step = optimizer.minimize(loss, var_list=opt_v)
num_epochs = graph_builders.get_or_create_num_epochs()
epoch_increment = tf.assign_add(num_epochs, 1)
train_ops = TrainOpsSupervised(
accumulate_gradients=None,
apply_gradients=train_step,
reset_gradients=None,
mc_step=mc_step,
acc_rate=acc_rate,
metrics=loss,
update_wf_norm=None,
epoch_increment=epoch_increment,
)
return train_ops
def build_opt_ops(
self,
wavefunction: wavefunctions.Wavefunction,
hamiltonian: operators.Operator,
hparams: tf.contrib.training.HParams,
shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
) -> NamedTuple:
"""Adds wavefunction optimization ops to the graph.
Args:
wavefunction: Wavefunction ansatz to optimize.
hamiltonian: Hamiltonian whose ground state we are solving for.
hparams: Hyperparameters of the optimization procedure.
shared_resources: Resources sharable among different modules.
Returns:
NamedTuple holding tensors needed to run a training epoch.
"""
batch_size = hparams.batch_size
n_sites = hparams.num_sites
configs = graph_builders.get_configs(shared_resources, batch_size, n_sites)
mc_step, acc_rate = graph_builders.get_monte_carlo_sampling(
shared_resources, configs, wavefunction)
opt_v = wavefunction.get_trainable_variables()
wf_omega = copy.deepcopy(wavefunction) # building supervisor wavefunction.
beta = tf.constant(hparams.time_evolution_beta, dtype=tf.float32)
beta2 = tf.constant(hparams.time_evolution_befta ** 2, dtype=tf.float32)
psi_omega = wf_omega(configs)
h_psi_omega = hamiltonian.apply_in_place(wf_omega, configs, psi_omega)
h_psi_omega_beta = h_psi_omega * beta
ite_psi_omega = psi_omega - h_psi_omega_beta
local_energy = h_psi_omega / psi_omega
energy_expectation = tf.reduce_mean(local_energy)
squared_energy_expectation = tf.reduce_mean(tf.square(local_energy))
ite_normalization = tf.sqrt(
1. - 2 * beta * energy_expectation + squared_energy_expectation * beta2
)
ite_normalization_var = tf.get_variable(
name='ite_normalization',
initializer=tf.ones(shape=[]),
dtype=tf.float32,
trainable=False
)
num_epochs = graph_builders.get_or_create_num_epochs()
exp_moving_average = tf.train.ExponentialMovingAverage(0.999, num_epochs)
accumulate_norm = exp_moving_average.apply([ite_normalization])
normalization_value = exp_moving_average.average(ite_normalization)
accumulate_energy = exp_moving_average.apply([energy_expectation])
energy_value = exp_moving_average.average(energy_expectation)
update_normalization = tf.assign(ite_normalization_var, normalization_value)
psi = wavefunction(configs)
loss = tf.reduce_mean(
tf.squared_difference(psi, ite_psi_omega / ite_normalization_var)
/ tf.square(tf.stop_gradient(psi))
)
optimizer = create_sgd_optimizer(hparams)
train_step = optimizer.minimize(loss, var_list=opt_v)
train_step = tf.group([train_step, accumulate_norm, accumulate_energy])
update_supervisor = wavefunctions.module_transfer_ops(
wavefunction, wf_omega)
epoch_increment = tf.assign_add(num_epochs, 1)
update_wf_norm = wavefunction.update_norm(psi)
train_ops = TrainOpsSWO(
train_step=train_step,
accumulate_gradients=None,
apply_gradients=None,
reset_gradients=None,
mc_step=mc_step,
acc_rate=acc_rate,
metrics=loss,
energy=energy_value,
update_supervisor=update_supervisor,
update_normalization=update_normalization,
epoch_increment=epoch_increment,
update_wf_norm=update_wf_norm,
)
return train_ops
def build_opt_ops(
self,
wavefunction: wavefunctions.Wavefunction,
hamiltonian: operators.Operator,
hparams: tf.contrib.training.HParams,
shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
) -> NamedTuple:
"""Adds wavefunction optimization ops to the graph.
Args:
wavefunction: Wavefunction ansatz to optimize.
hamiltonian: Hamiltonian whose ground state we are solving for.
hparams: Hyperparameters of the optimization procedure.
shared_resources: Resources sharable among different modules.
Returns:
NamedTuple holding tensors needed to run a training epoch.
"""
batch_size = hparams.batch_size
n_sites = hparams.num_sites
configs = graph_builders.get_configs(shared_resources, batch_size, n_sites)
mc_step, acc_rate = graph_builders.get_monte_carlo_sampling(
shared_resources, configs, wavefunction)
opt_v = wavefunction.get_trainable_variables()
wf_omega = copy.deepcopy(wavefunction) # building supervisor wavefunction.
psi = wavefunction(configs)
psi_omega = wf_omega(configs)
beta = tf.constant(hparams.time_evolution_beta, dtype=tf.float32)
h_psi_omega = hamiltonian.apply_in_place(wf_omega, configs, psi_omega)
h_psi_omega_beta = beta * h_psi_omega
ite_psi_omega = psi_omega - h_psi_omega_beta
local_energy = h_psi_omega / psi_omega
update_wf_norm = wavefunction.update_norm(psi)
psi_no_grad = tf.stop_gradient(psi)
ratio = tf.stop_gradient(ite_psi_omega / psi)
log_psi_raw_grads = tf.gradients(psi / psi_no_grad, opt_v)
log_psi_grads = [tf.convert_to_tensor(grad) for grad in log_psi_raw_grads]
ratio_log_psi_raw_grads = tf.gradients(ratio * psi / psi_no_grad, opt_v)
ratio_log_psi_grads = [
tf.convert_to_tensor(grad) for grad in ratio_log_psi_raw_grads
]
log_grads = [
tf.metrics.mean_tensor(log_psi_grad) for log_psi_grad in log_psi_grads
]
ratio_grads = [
tf.metrics.mean_tensor(grad) for grad in ratio_log_psi_grads
]
mean_energy, accumulate_energy = tf.metrics.mean(local_energy)
mean_ratio, accumulate_ratio = tf.metrics.mean(ratio)
mean_log_grads, accumulate_log_grads = list(map(list, zip(*log_grads)))
mean_ratio_grads, accumulate_ratio_grads = list(
map(list, zip(*ratio_grads)))
grad_pairs = zip(mean_log_grads, mean_ratio_grads)
overlap_gradients = [
grad - scaled_grad / mean_ratio for grad, scaled_grad in grad_pairs
]
grads_and_vars = list(zip(overlap_gradients, opt_v))
optimizer = create_sgd_optimizer(hparams)
apply_gradients = optimizer.apply_gradients(grads_and_vars)
all_updates = [accumulate_ratio, accumulate_energy]
all_updates += accumulate_log_grads + accumulate_ratio_grads
accumulate_gradients = tf.group(all_updates)
update_network = wavefunctions.module_transfer_ops(wavefunction, wf_omega)
clear_gradients = tf.variables_initializer(tf.local_variables())
num_epochs = graph_builders.get_or_create_num_epochs()
epoch_increment = tf.assign_add(num_epochs, 1)
train_ops = TrainOpsSWO(
train_step=None,
accumulate_gradients=accumulate_gradients,
apply_gradients=apply_gradients,
reset_gradients=clear_gradients,
mc_step=mc_step,
acc_rate=acc_rate,
metrics=None,
energy=mean_energy,
update_supervisor=update_network,
update_normalization=None,
epoch_increment=epoch_increment,
update_wf_norm=update_wf_norm,
)
return train_ops
def build_opt_ops(
self,
wavefunction: wavefunctions.Wavefunction,
hamiltonian: operators.Operator,
hparams: tf.contrib.training.HParams,
shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
) -> NamedTuple:
"""Adds wavefunction optimization ops to the graph.
Args:
wavefunction: Wavefunction ansatz to optimize.
hamiltonian: Hamiltonian whose ground state we are solving for.
hparams: Hyperparameters of the optimization procedure.
shared_resources: Resources sharable among different modules.
Returns:
NamedTuple holding tensors needed to run a training epoch.
"""
batch_size = hparams.batch_size
n_sites = hparams.num_sites
configs = graph_builders.get_configs(shared_resources, batch_size, n_sites)
mc_step, acc_rate = graph_builders.get_monte_carlo_sampling(
shared_resources, configs, wavefunction)
opt_v = wavefunction.get_trainable_variables()
psi = wavefunction(configs)
psi_no_grad = tf.stop_gradient(psi)
update_wf_norm = wavefunction.update_norm(psi)
local_energy = hamiltonian.local_value(wavefunction, configs, psi)
local_energy = tf.stop_gradient(local_energy)
log_psi_raw_grads = tf.gradients(psi / psi_no_grad, opt_v)
log_psi_grads = [tf.convert_to_tensor(grad) for grad in log_psi_raw_grads]
e_psi_raw_grads = tf.gradients(psi / psi_no_grad * local_energy, opt_v)
e_psi_grads = [tf.convert_to_tensor(grad) for grad in e_psi_raw_grads]
grads = [
tf.metrics.mean_tensor(log_psi_grad) for log_psi_grad in log_psi_grads
]
weighted_grads = [tf.metrics.mean_tensor(grad) for grad in e_psi_grads]
mean_energy, update_energy = tf.metrics.mean(local_energy)
mean_pure_grads, update_pure_grads = list(map(list, zip(*grads)))
mean_scaled_grads, update_scaled_grads = list(
map(list, zip(*weighted_grads)))
grad_pairs = zip(mean_pure_grads, mean_scaled_grads)
energy_gradients = [
scaled_grad - mean_energy * grad for grad, scaled_grad in grad_pairs
]
grads_and_vars = list(zip(energy_gradients, opt_v))
optimizer = create_sgd_optimizer(hparams)
apply_gradients = optimizer.apply_gradients(grads_and_vars)
reset_gradients = tf.variables_initializer(tf.local_variables())
all_updates = [update_energy,] + update_pure_grads + update_scaled_grads
accumulate_gradients = tf.group(all_updates)
num_epochs = graph_builders.get_or_create_num_epochs()
epoch_increment = tf.assign_add(num_epochs, 1)
train_ops = TrainOpsTraditional(
accumulate_gradients=accumulate_gradients,
apply_gradients=apply_gradients,
reset_gradients=reset_gradients,
mc_step=mc_step,
acc_rate=acc_rate,
metrics=mean_energy,
epoch_increment=epoch_increment,
update_wf_norm=update_wf_norm,
)
return train_ops
def build_opt_ops(
self,
wavefunction: wavefunctions.Wavefunction,
target_wavefunction: wavefunctions.Wavefunction,
hparams: tf.contrib.training.HParams,
shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
) -> NamedTuple:
"""Adds wavefunction optimization ops to the graph.
Args:
wavefunction: Wavefunction model to optimize.
target_wavefunction: Wavefunction model that serves as a target state.
shared_resources: Resources sharable among different modules.
hparams: Hyperparameters of the optimization procedure.
Returns:
NamedTuple holding tensors needed to run a training epoch.
"""
batch_size = hparams.batch_size
n_sites = hparams.num_sites
psi_configs = graph_builders.get_configs(
shared_resources, batch_size // 2, n_sites)
target_configs = graph_builders.get_configs(
shared_resources, batch_size // 2, n_sites,
configs_id=graph_builders.ResourceName.TARGET_CONFIGS)
# Here we use explicit call to build separate MCMC sampling for two wf.
target_mc_step, target_acc_rate = graph_builders.build_monte_carlo_sampling(
target_configs, target_wavefunction)
psi_mc_step, psi_acc_rate = graph_builders.build_monte_carlo_sampling(
psi_configs, wavefunction)
mc_step = tf.group([psi_mc_step, target_mc_step])
acc_rate = tf.group([psi_acc_rate, target_acc_rate])
configs = tf.concat([psi_configs, target_configs], axis=0)
psi = wavefunction(configs)
psi_target = target_wavefunction(configs) * np.sqrt(2 ** n_sites)
# # A version of accounting for sampling bias.
# psi_no_grad = tf.stop_gradient(psi)
# loss = tf.reduce_mean(
# tf.squared_difference(psi, psi_target) /
# (tf.square(psi_no_grad) + tf.square(psi_target))
# )
loss = tf.reduce_mean(
tf.squared_difference(psi, psi_target)
)
opt_v = wavefunction.get_trainable_variables()
optimizer = create_sgd_optimizer(hparams)
train_step = optimizer.minimize(loss, var_list=opt_v)
num_epochs = graph_builders.get_or_create_num_epochs()
epoch_increment = tf.assign_add(num_epochs, 1)
train_ops = TrainOpsSupervised(
accumulate_gradients=None,
apply_gradients=train_step,
reset_gradients=None,
mc_step=mc_step,
acc_rate=acc_rate,
metrics=loss,
update_wf_norm=None,
epoch_increment=epoch_increment,
)
return train_ops
def build_opt_ops(
self,
wavefunction: wavefunctions.Wavefunction,
target_wavefunction: wavefunctions.Wavefunction,
hparams: tf.contrib.training.HParams,
shared_resources: Dict[graph_builders.ResourceName, tf.Tensor],
) -> NamedTuple:
"""Adds wavefunction optimization ops to the graph.
Args:
wavefunction: Wavefunction model to optimize.
target_wavefunction: Wavefunction model that serves as a target state.
shared_resources: Resources sharable among different modules.
hparams: Hyperparameters of the optimization procedure.
Returns:
NamedTuple holding tensors needed to run a training epoch.
"""
batch_size = hparams.batch_size
n_sites = hparams.num_sites
configs = graph_builders.get_configs(shared_resources, batch_size, n_sites)
mc_step, acc_rate = graph_builders.get_monte_carlo_sampling(
shared_resources, configs, wavefunction)
psi = wavefunction(configs)
psi_target = target_wavefunction(configs)
psi_no_grad = tf.stop_gradient(psi)
opt_v = wavefunction.get_trainable_variables()
# Computation of L=log(overlap) gradient with respect to opt_v using
# "grad(L) = <grad(log(psi))> - <grad(log(psi))*ratio> / <ratio>"
# where ratio is the ratio of target and curren wave-function.
ratio = tf.stop_gradient(psi_target / psi)
log_psi_raw_grads = tf.gradients(psi / psi_no_grad, opt_v)
log_psi_grads = [tf.convert_to_tensor(grad) for grad in log_psi_raw_grads]
ratio_log_psi_raw_grads = tf.gradients(ratio * psi / psi_no_grad, opt_v)
ratio_log_psi_grads = [
tf.convert_to_tensor(grad) for grad in ratio_log_psi_raw_grads
]
log_grads = [
tf.metrics.mean_tensor(log_psi_grad) for log_psi_grad in log_psi_grads
]
ratio_grads = [
tf.metrics.mean_tensor(grad) for grad in ratio_log_psi_grads
]
mean_ratio, accumulate_ratio = tf.metrics.mean(ratio)
mean_log_grads, accumulate_log_grads = list(map(list, zip(*log_grads)))
mean_ratio_grads, accumulate_ratio_grads = list(
map(list, zip(*ratio_grads)))
all_updates = accumulate_log_grads + accumulate_ratio_grads
all_updates.append(accumulate_ratio)
accumulate_gradients = tf.group(all_updates)
optimizer = create_sgd_optimizer(hparams)
grad_pairs = zip(mean_log_grads, mean_ratio_grads)
overlap_gradients = [
grad - scaled_grad / mean_ratio for grad, scaled_grad in grad_pairs
]
grads_and_vars = list(zip(overlap_gradients, opt_v))
apply_gradients = optimizer.apply_gradients(grads_and_vars)
reset_gradients = tf.variables_initializer(tf.local_variables())
update_wf_norm = wavefunction.update_norm(psi)
num_epochs = graph_builders.get_or_create_num_epochs()
epoch_increment = tf.assign_add(num_epochs, 1)
train_ops = TrainOpsSupervised(
accumulate_gradients=accumulate_gradients,
apply_gradients=apply_gradients,
reset_gradients=reset_gradients,
mc_step=mc_step,
acc_rate=acc_rate,
metrics=None,
update_wf_norm=update_wf_norm,
epoch_increment=epoch_increment,
)
return train_ops