optimizer = Adam(lr=0.001, beta_1=0.9, beta_2=0.999)
model.compile(optimizer=optimizer, loss=loss_for_energy_minimization)
model.summary()
operator = Ising(h=3.0, hilbert_state_shape=hilbert_state_shape, pbc=False)
sampler = AutoregressiveSampler(conditional_log_probs_model, batch_size)
monte_carlo_generator = VariationalMonteCarlo(model, operator, sampler)
validation_sampler = AutoregressiveSampler(conditional_log_probs_model,
batch_size * 16)
validation_generator = VariationalMonteCarlo(model, operator,
validation_sampler)
tensorboard = TensorBoardWithGeneratorValidationData(
log_dir='tensorboard_logs/invariant_example_2d_monte_carlo_batch_%s_run_3'
% batch_size,
generator=monte_carlo_generator,
update_freq=1,
histogram_freq=1,
batch_size=batch_size,
write_output=False)
callbacks = default_wave_function_stats_callbacks_factory(
monte_carlo_generator,
validation_generator=validation_generator,
true_ground_state_energy=-50.18662388277671) + [tensorboard]
model.fit_generator(monte_carlo_generator,
steps_per_epoch=steps_per_epoch,
epochs=2,
callbacks=callbacks,
max_queue_size=0,
workers=0)
model.save_weights('final_2d_ising_fcnn.h5')
loss=loss_for_energy_minimization,
metrics=optimizer.metrics)
model.summary()
operator = Heisenberg(hilbert_state_shape=hilbert_state_shape, pbc=True)
sampler = MetropolisHastingsHamiltonian(
model,
batch_size,
operator,
num_of_chains=20,
unused_sampels=numpy.prod(hilbert_state_shape))
variational_monte_carlo = VariationalMonteCarlo(model, operator, sampler)
tensorboard = TensorBoardWithGeneratorValidationData(
log_dir='tensorboard_logs/rbm_with_sr_run_6',
generator=variational_monte_carlo,
update_freq=1,
histogram_freq=1,
batch_size=batch_size,
write_output=False)
callbacks = default_wave_function_stats_callbacks_factory(
variational_monte_carlo, true_ground_state_energy=-35.6175461195) + [
MCMCStats(variational_monte_carlo), tensorboard
]
model.fit_generator(variational_monte_carlo.to_generator(),
steps_per_epoch=steps_per_epoch,
epochs=1,
callbacks=callbacks,
max_queue_size=0,
workers=0)
model.save_weights('final_1d_heisenberg.h5')
def train(operator, config, true_ground_state_energy=None):
if config.use_horovod:
init_horovod()
to_valid_stages_config(config)
is_rank_0 = (not config.use_horovod) or hvd.rank() == 0
if is_rank_0:
save_config(config)
model, sampler = build_model(operator, config)
optimizer = compile_model(model, config.learning_rate[0],
config.use_horovod)
checkpoint_path = os.path.join(config.output_path, 'model.h5')
initial_epoch = load_weights_if_exist(model, checkpoint_path)
total_epochs = 0
for idx, (batch_size, num_epoch, learning_rate) in enumerate(
zip(config.batch_size, config.num_epoch, config.learning_rate)):
total_epochs += num_epoch
if total_epochs <= initial_epoch:
continue
vmc_cls = VariationalMonteCarlo
if config.use_horovod:
batch_size = int(math.ceil(batch_size / hvd.size()))
vmc_cls = HorovodVariationalMonteCarlo
validation_sampler = sampler.copy_with_new_batch_size(
min(batch_size * 8, 2**15), mini_batch_size=config.mini_batch_size)
assert batch_size < config.mini_batch_size or batch_size % config.mini_batch_size == 0
sampler = sampler.copy_with_new_batch_size(batch_size,
config.mini_batch_size)
variational_monte_carlo = vmc_cls(
model, operator, sampler, mini_batch_size=config.mini_batch_size)
validation_generator = vmc_cls(
model,
operator,
validation_sampler,
wave_function_evaluation_batch_size=config.mini_batch_size)
optimizer.set_update_params_frequency(
variational_monte_carlo.update_params_frequency)
K.set_value(optimizer.lr, learning_rate)
callbacks = default_wave_function_stats_callbacks_factory(
variational_monte_carlo,
validation_generator=validation_generator,
log_in_batch_or_epoch=False,
true_ground_state_energy=true_ground_state_energy,
validation_period=config.validation_period)
if config.use_horovod:
callbacks = [
hvd.callbacks.BroadcastGlobalVariablesCallback(0),
] + callbacks + [hvd.callbacks.MetricAverageCallback()]
if is_rank_0:
tensorboard = TensorBoardWithGeneratorValidationData(
log_dir=config.output_path,
generator=variational_monte_carlo,
update_freq='epoch',
histogram_freq=0,
batch_size=batch_size,
write_output=False,
write_graph=False)
callbacks += [
tensorboard,
CheckpointByTime(checkpoint_path, save_weights_only=True)
]
verbose = 1
else:
verbose = 0
model.fit_generator(variational_monte_carlo.to_generator(),
steps_per_epoch=config.steps_per_epoch *
variational_monte_carlo.update_params_frequency,
epochs=total_epochs,
callbacks=callbacks,
max_queue_size=0,
workers=0,
initial_epoch=initial_epoch,
verbose=verbose)
if is_rank_0:
model.save_weights(
os.path.join(config.output_path, 'stage_%s.h5' % (idx + 1)))
initial_epoch = total_epochs
evaluation_inputs = Input(shape=config.hilbert_state_shape, dtype='int8')
obc_input = Input(shape=config.hilbert_state_shape,
dtype=evaluation_inputs.dtype)
invariant_model = make_2d_obc_invariants(obc_input, model)
invariant_model = make_up_down_invariant(evaluation_inputs,
invariant_model)
mini_batch_size = config.mini_batch_size // 16
sampler = sampler.copy_with_new_batch_size(config.mini_batch_size)
vmc_cls = VariationalMonteCarlo
if config.use_horovod:
vmc_cls = HorovodVariationalMonteCarlo
variational_monte_carlo = vmc_cls(invariant_model,
operator,
sampler,
mini_batch_size=config.mini_batch_size)
callbacks = default_wave_function_stats_callbacks_factory(
variational_monte_carlo,
log_in_batch_or_epoch=False,
true_ground_state_energy=true_ground_state_energy)
if config.use_horovod:
callbacks = callbacks + [hvd.callbacks.MetricAverageCallback()]
results = evaluate(variational_monte_carlo,
steps=(2**15) // mini_batch_size,
callbacks=callbacks[:4],
keys_to_progress_bar_mapping={
'energy/energy': 'energy',
'energy/relative_error': 'relative_error',
'energy/local_energy_variance': 'variance'
},
verbose=is_rank_0)
if is_rank_0:
print(results)
def run(params, batch_size_list, epochs_list):
run_name = 'depth_%s_width_%s_weights_normalization_%s_adam_lr_%s_gamma_%s_run_%s' % (
params.depth, params.width, params.no_weights_normalization,
params.learning_rate, params.gamma, params.run_index)
hilbert_state_shape = (12, 12)
model, sampler = build_model(hilbert_state_shape, params.depth,
params.width,
not params.no_weights_normalization,
params.learning_rate)
operator = Ising(hilbert_state_shape=hilbert_state_shape,
pbc=False,
h=params.gamma)
checkpoint_path = '%s.h5' % run_name
initial_epoch = restore_run_state(model, checkpoint_path)
total_epochs = 0
mini_batch_size = depth_to_max_mini_batch(params.depth)
true_ground_state_energy = true_ground_state_energy_mapping[params.gamma]
for idx, (batch_size,
epochs) in enumerate(zip(batch_size_list, epochs_list)):
total_epochs += epochs
if total_epochs <= initial_epoch:
continue
validation_sampler = sampler.copy_with_new_batch_size(
min(batch_size * 8, 2**15), mini_batch_size=mini_batch_size)
assert batch_size < mini_batch_size or batch_size % mini_batch_size == 0
sampler = sampler.copy_with_new_batch_size(batch_size, mini_batch_size)
variational_monte_carlo = VariationalMonteCarlo(
model, operator, sampler, mini_batch_size=mini_batch_size)
validation_generator = VariationalMonteCarlo(
model,
operator,
validation_sampler,
wave_function_evaluation_batch_size=mini_batch_size)
model.optimizer.set_update_params_frequency(
variational_monte_carlo.update_params_frequency)
tensorboard = TensorBoardWithGeneratorValidationData(
log_dir='tensorboard_logs/%s' % run_name,
generator=variational_monte_carlo,
update_freq='epoch',
histogram_freq=0,
batch_size=batch_size,
write_output=False)
callbacks = default_wave_function_stats_callbacks_factory(
variational_monte_carlo,
validation_generator=validation_generator,
log_in_batch_or_epoch=False,
true_ground_state_energy=true_ground_state_energy,
validation_period=3)
callbacks += [
tensorboard,
CheckpointByTime(checkpoint_path, save_weights_only=True)
]
model.fit_generator(variational_monte_carlo.to_generator(),
steps_per_epoch=100 *
variational_monte_carlo.update_params_frequency,
epochs=total_epochs,
callbacks=callbacks,
max_queue_size=0,
workers=0,
initial_epoch=initial_epoch)
model.save_weights('%s_stage_%s.h5' % (run_name, idx))
initial_epoch = total_epochs
evaluation_inputs = Input(shape=hilbert_state_shape, dtype='int8')
obc_input = Input(shape=hilbert_state_shape, dtype=evaluation_inputs.dtype)
invariant_model = make_2d_obc_invariants(obc_input, model)
invariant_model = make_up_down_invariant(evaluation_inputs,
invariant_model)
mini_batch_size = mini_batch_size // 16
sampler = sampler.copy_with_new_batch_size(mini_batch_size)
variational_monte_carlo = VariationalMonteCarlo(
invariant_model, operator, sampler, mini_batch_size=mini_batch_size)
callbacks = default_wave_function_stats_callbacks_factory(
variational_monte_carlo,
log_in_batch_or_epoch=False,
true_ground_state_energy=true_ground_state_energy)
results = evaluate(variational_monte_carlo,
steps=(2**15) // mini_batch_size,
callbacks=callbacks[:4],
keys_to_progress_bar_mapping={
'energy/energy': 'energy',
'energy/relative_error': 'relative_error',
'energy/local_energy_variance': 'variance'
})
print(results)
# sampler = FastAutoregressiveSampler(fast_sampling, buffer_size=5000)
# sampler = FastAutoregressiveSampler(conditional_log_probs)
generator = VariationalMonteCarlo(model,
operator,
sampler,
cache=wave_function_cache)
#### Exact Grads ####
# generator = ExactVariational(model, operator, batch_size, cache=wave_function_cache)
checkpoint = ModelCheckpoint('ising_fcnn.h5',
monitor='energy',
save_best_only=True,
save_weights_only=True)
# tensorboard = TensorBoard(update_freq=1)
tensorboard = TensorBoardWithGeneratorValidationData(
monte_carlo_iterator=generator,
update_freq=1,
histogram_freq=1,
batch_size=batch_size)
early_stopping = EarlyStopping(monitor='relative_energy_error', min_delta=1e-5)
callbacks = [
LocalEnergyStats(generator, true_ground_state_energy=-457.0416241),
wave_function_cache,
LocalStats("Energy Again",
sampler=sampler,
operator=operator,
cache=wave_function_cache),
SigmaZStats(monte_carlo_generator=generator), checkpoint, tensorboard,
early_stopping,
TerminateOnNaN()
]
model.fit_generator(generator,