def __call__(self, iteration, theta_loss, phi_loss, generative_model,
inference_network, memory, optimizer):
if iteration % self.logging_interval == 0:
util.print_with_time(
'Iteration {} losses: theta = {:.3f}, phi = {:.3f}'.format(
iteration, theta_loss, phi_loss))
self.theta_loss_history.append(theta_loss)
self.phi_loss_history.append(phi_loss)
if iteration % self.checkpoint_interval == 0:
stats_filename = util.get_stats_path(self.model_folder)
util.save_object(self, stats_filename)
util.save_models(generative_model, inference_network,
self.model_folder, iteration, memory)
if iteration % self.eval_interval == 0:
self.p_error_history.append(
util.get_p_error(self.true_generative_model, generative_model))
self.q_error_history.append(
util.get_q_error(self.true_generative_model, inference_network,
self.test_obss))
# TODO
# self.memory_error_history.append(util.get_memory_error(
# self.true_generative_model, memory, generative_model,
# self.test_obss))
util.print_with_time(
'Iteration {} p_error = {:.3f}, q_error_to_true = '
'{:.3f}'.format(iteration, self.p_error_history[-1],
self.q_error_history[-1]))
def __call__(self, iteration, loss, elbo, generative_model,
inference_network, control_variate):
if iteration % self.logging_interval == 0:
util.print_with_time(
'Iteration {} loss = {:.3f}, elbo = {:.3f}'.format(
iteration, loss, elbo))
self.loss_history.append(loss)
self.elbo_history.append(elbo)
if iteration % self.checkpoint_interval == 0:
stats_filename = util.get_stats_filename(self.model_folder)
util.save_object(self, stats_filename)
util.save_models(generative_model, inference_network,
self.pcfg_path, self.model_folder)
util.save_control_variate(control_variate, self.model_folder)
if iteration % self.eval_interval == 0:
self.p_error_history.append(
util.get_p_error(self.true_generative_model, generative_model))
self.q_error_to_true_history.append(
util.get_q_error(self.true_generative_model,
inference_network))
self.q_error_to_model_history.append(
util.get_q_error(generative_model, inference_network))
util.print_with_time(
'Iteration {} p_error = {:.3f}, q_error_to_true = {:.3f}, '
'q_error_to_model = {:.3f}'.format(
iteration, self.p_error_history[-1],
self.q_error_to_true_history[-1],
self.q_error_to_model_history[-1]))
def __call__(self, iteration, wake_theta_loss, wake_phi_loss, elbo,
generative_model, inference_network, optimizer_theta,
optimizer_phi):
if iteration % self.logging_interval == 0:
util.print_with_time(
'Iteration {} losses: theta = {:.3f}, phi = {:.3f}, elbo = '
'{:.3f}'.format(iteration, wake_theta_loss, wake_phi_loss,
elbo))
self.wake_theta_loss_history.append(wake_theta_loss)
self.wake_phi_loss_history.append(wake_phi_loss)
self.elbo_history.append(elbo)
if iteration % self.checkpoint_interval == 0:
stats_filename = util.get_stats_filename(self.model_folder)
util.save_object(self, stats_filename)
util.save_models(generative_model, inference_network,
self.pcfg_path, self.model_folder)
if iteration % self.eval_interval == 0:
self.p_error_history.append(
util.get_p_error(self.true_generative_model, generative_model))
self.q_error_to_true_history.append(
util.get_q_error(self.true_generative_model,
inference_network))
self.q_error_to_model_history.append(
util.get_q_error(generative_model, inference_network))
util.print_with_time(
'Iteration {} p_error = {:.3f}, q_error_to_true = {:.3f}, '
'q_error_to_model = {:.3f}'.format(
iteration, self.p_error_history[-1],
self.q_error_to_true_history[-1],
self.q_error_to_model_history[-1]))
def main():
batch_size = 2
pcfg_path = './pcfgs/astronomers_pcfg.json'
generative_model, inference_network, true_generative_model = \
util.init_models(pcfg_path)
obss = [true_generative_model.sample_obs() for _ in range(batch_size)]
num_mc_samples = 100
num_particles_list = [2, 5, 10, 20, 50, 100]
vimco_grad = np.zeros((len(num_particles_list), 2))
vimco_one_grad = np.zeros((len(num_particles_list), 2))
reinforce_grad = np.zeros((len(num_particles_list), 2))
reinforce_one_grad = np.zeros((len(num_particles_list), 2))
two_grad = np.zeros((len(num_particles_list), 2))
log_evidence_stats = np.zeros((len(num_particles_list), 2))
log_evidence_grad = np.zeros((len(num_particles_list), 2))
wake_phi_loss_grad = np.zeros((len(num_particles_list), 2))
log_Q_grad = np.zeros((len(num_particles_list), 2))
sleep_loss_grad = np.zeros((len(num_particles_list), 2))
for i, num_particles in enumerate(num_particles_list):
util.print_with_time('num_particles = {}'.format(num_particles))
(vimco_grad[i], vimco_one_grad[i], reinforce_grad[i],
reinforce_one_grad[i], two_grad[i], log_evidence_stats[i],
log_evidence_grad[i], wake_phi_loss_grad[i], log_Q_grad[i],
sleep_loss_grad[i]) = get_mean_stds(
generative_model, inference_network, num_mc_samples, obss,
num_particles)
util.save_object([
vimco_grad, vimco_one_grad, reinforce_grad, reinforce_one_grad,
two_grad, log_evidence_stats, log_evidence_grad, wake_phi_loss_grad,
log_Q_grad, sleep_loss_grad],
'./variance_analysis/data.pkl')
def run(args):
util.print_with_time(str(args))
# save args
model_folder = util.get_model_folder()
args_filename = util.get_args_filename(model_folder)
util.save_object(args, args_filename)
# init models
util.set_seed(args.seed)
generative_model, inference_network, true_generative_model = \
load_or_init_models(args.load_model_folder, args.pcfg_path)
if args.train_mode == 'relax':
control_variate = models.ControlVariate(generative_model.grammar)
# train
if args.train_mode == 'ww':
train_callback = train.TrainWakeWakeCallback(
args.pcfg_path, model_folder, true_generative_model,
args.logging_interval, args.checkpoint_interval,
args.eval_interval)
train.train_wake_wake(generative_model, inference_network,
true_generative_model, args.batch_size,
args.num_iterations, args.num_particles,
train_callback)
elif args.train_mode == 'ws':
train_callback = train.TrainWakeSleepCallback(
args.pcfg_path, model_folder, true_generative_model,
args.logging_interval, args.checkpoint_interval,
args.eval_interval)
train.train_wake_sleep(generative_model, inference_network,
true_generative_model, args.batch_size,
args.num_iterations, args.num_particles,
train_callback)
elif args.train_mode == 'reinforce' or args.train_mode == 'vimco':
train_callback = train.TrainIwaeCallback(
args.pcfg_path, model_folder, true_generative_model,
args.logging_interval, args.checkpoint_interval,
args.eval_interval)
train.train_iwae(args.train_mode, generative_model, inference_network,
true_generative_model, args.batch_size,
args.num_iterations, args.num_particles,
train_callback)
elif args.train_mode == 'relax':
train_callback = train.TrainRelaxCallback(
args.pcfg_path, model_folder, true_generative_model,
args.logging_interval, args.checkpoint_interval,
args.eval_interval)
train.train_relax(generative_model, inference_network, control_variate,
true_generative_model, args.batch_size,
args.num_iterations, args.num_particles,
train_callback)
# save models and stats
util.save_models(generative_model, inference_network, args.pcfg_path,
model_folder)
stats_filename = util.get_stats_filename(model_folder)
util.save_object(train_callback, stats_filename)
def _train_and_predict_bcse(study_id, train_data, competency_attempt_counters, predictions):
try:
predictions.extend(BasicColdStartMasteryEstimator().predict(train_data))
for i, pred in enumerate(predictions):
learner_id = pred['learnerId']
competency_id = inference_utils.sanitize_competency_id(pred['competencyId'])
predictions[i]['nextAttempt'] = competency_attempt_counters[learner_id][competency_id] + 1
predictions[i]['@type'] = 'MasteryProbability'
except Exception as e:
print_with_time('ERROR: ' + traceback.format_exc())
def train_and_predict(study_id, test_statements):
train_data = CompetencyAttemptsCollector().get_competency_attempts(test_statements)
predictions_afm = list()
predictions_transfer_learning = list()
predictions_bcse = list()
if len(train_data) > 0:
# Populate competency_attempt_counters in advance.
competency_attempt_counters = {}
for learner_id, competency_dict in train_data.items():
# Need to use list to be 100% sure we're not modifying the dict while iterating
for competency_id in list(competency_dict.keys()):
competency_dict[inference_utils.sanitize_competency_id(competency_id)] = competency_dict.pop(competency_id)
if learner_id not in competency_attempt_counters:
competency_attempt_counters[learner_id] = {}
for competency_id in competency_dict.keys():
try:
attempt_counter = get_competency_attempt_counters(learner_id, competency_id)
except Exception as e:
print_with_time('ERROR: ' + traceback.format_exc())
attempt_counter = None
if attempt_counter is None:
competency_attempt_counters[learner_id][competency_id] = 0
else:
competency_attempt_counters[learner_id][competency_id] = attempt_counter['numAttempts']
if AFM_ENABLED:
afm_thread = Thread(target=_train_and_predict_afm,
args=(study_id, train_data, competency_attempt_counters, predictions_afm))
afm_thread.start()
if TRANSFER_LEARNING_ENABLED:
transfer_learning_thread = Thread(target=_train_and_predict_transfer_learning,
args=(study_id, train_data, competency_attempt_counters, predictions_transfer_learning))
transfer_learning_thread.start()
if BCSE_ENABLED:
bcse_thread = Thread(target=_train_and_predict_bcse,
args=(study_id, train_data, competency_attempt_counters, predictions_bcse))
bcse_thread.start()
if AFM_ENABLED:
afm_thread.join()
if TRANSFER_LEARNING_ENABLED:
transfer_learning_thread.join()
if BCSE_ENABLED:
bcse_thread.join()
return predictions_afm + predictions_transfer_learning + predictions_bcse
def __call__(self, iteration, wake_theta_loss, wake_phi_loss, elbo,
generative_model, inference_network, optimizer_theta,
optimizer_phi):
if iteration % self.logging_interval == 0:
util.print_with_time(
'Iteration {} losses: theta = {:.3f}, phi = {:.3f}, elbo = '
'{:.3f}'.format(iteration, wake_theta_loss, wake_phi_loss,
elbo))
self.wake_theta_loss_history.append(wake_theta_loss)
self.wake_phi_loss_history.append(wake_phi_loss)
self.elbo_history.append(elbo)
if iteration % self.checkpoint_interval == 0:
stats_path = util.get_stats_path(self.save_dir)
util.save_object(self, stats_path)
util.save_checkpoint(self.save_dir,
iteration,
generative_model=generative_model,
inference_network=inference_network)
if iteration % self.eval_interval == 0:
log_p, kl = eval_gen_inf(generative_model, inference_network,
self.test_data_loader,
self.eval_num_particles)
self.log_p_history.append(log_p)
self.kl_history.append(kl)
stats = util.OnlineMeanStd()
for _ in range(10):
generative_model.zero_grad()
wake_theta_loss, elbo = losses.get_wake_theta_loss(
generative_model, inference_network, self.test_obs,
self.num_particles)
wake_theta_loss.backward()
theta_grads = [
p.grad.clone() for p in generative_model.parameters()
]
inference_network.zero_grad()
wake_phi_loss = losses.get_wake_phi_loss(
generative_model, inference_network, self.test_obs,
self.num_particles)
wake_phi_loss.backward()
phi_grads = [p.grad for p in inference_network.parameters()]
stats.update(theta_grads + phi_grads)
self.grad_std_history.append(stats.avg_of_means_stds()[1].item())
util.print_with_time(
'Iteration {} log_p = {:.3f}, kl = {:.3f}'.format(
iteration, self.log_p_history[-1], self.kl_history[-1]))
def _train_and_predict_transfer_learning(study_id, train_data, competency_attempt_counters, predictions):
# TODO: To improve performance, load the models at service startup and keep them in memory.
# Storing models after each train is still needed to cover for potential system crashes.
try:
online_transfer, online_transfer_file = _retrieve_online_transfer_estimator(study_id)
online_transfer.train(train_data)
predictions.extend(online_transfer.predict(train_data))
for i, pred in enumerate(predictions):
learner_id = pred['learnerId']
competency_id = inference_utils.sanitize_competency_id(pred['competencyId'])
predictions[i]['nextAttempt'] = competency_attempt_counters[learner_id][competency_id] + 1
predictions[i]['@type'] = 'MasteryProbability'
pickle.dump(online_transfer, open(online_transfer_file, "wb"))
except Exception as e:
print_with_time('ERROR: ' + traceback.format_exc())
def main():
num_mixtures = 20
temp = np.arange(num_mixtures) + 5
true_p_mixture_probs = temp / np.sum(temp)
softmax_multiplier = 0.5
args = argparse.Namespace(
init_mixture_logits=np.array(
list(reversed(2 * np.arange(num_mixtures)))),
softmax_multiplier=softmax_multiplier,
device=torch.device('cpu'),
num_mixtures=num_mixtures,
relaxed_one_hot=False,
temperature=None,
true_mixture_logits=np.log(true_p_mixture_probs) / softmax_multiplier)
batch_size = 2
generative_model, inference_network, true_generative_model = \
util.init_models(args)
obs = true_generative_model.sample_obs(batch_size)
num_mc_samples = 100
num_particles_list = [2, 5, 10, 20, 50, 100]
vimco_grad = np.zeros((len(num_particles_list), 2))
vimco_one_grad = np.zeros((len(num_particles_list), 2))
reinforce_grad = np.zeros((len(num_particles_list), 2))
reinforce_one_grad = np.zeros((len(num_particles_list), 2))
two_grad = np.zeros((len(num_particles_list), 2))
log_evidence_stats = np.zeros((len(num_particles_list), 2))
log_evidence_grad = np.zeros((len(num_particles_list), 2))
wake_phi_loss_grad = np.zeros((len(num_particles_list), 2))
log_Q_grad = np.zeros((len(num_particles_list), 2))
sleep_loss_grad = np.zeros((len(num_particles_list), 2))
for i, num_particles in enumerate(num_particles_list):
util.print_with_time('num_particles = {}'.format(num_particles))
(vimco_grad[i], vimco_one_grad[i], reinforce_grad[i],
reinforce_one_grad[i], two_grad[i], log_evidence_stats[i],
log_evidence_grad[i], wake_phi_loss_grad[i], log_Q_grad[i],
sleep_loss_grad[i]) = get_mean_stds(generative_model,
inference_network, num_mc_samples,
obs, num_particles)
util.save_object([
vimco_grad, vimco_one_grad, reinforce_grad, reinforce_one_grad,
two_grad, log_evidence_stats, log_evidence_grad, wake_phi_loss_grad,
log_Q_grad, sleep_loss_grad
], './variance_analysis/data.pkl')
def __call__(self, iteration, loss, elbo, generative_model,
inference_network, optimizer):
if iteration % self.logging_interval == 0:
util.print_with_time(
'Iteration {} loss = {:.3f}, elbo = {:.3f}'.format(
iteration, loss, elbo))
self.loss_history.append(loss)
self.elbo_history.append(elbo)
if iteration % self.checkpoint_interval == 0:
stats_path = util.get_stats_path(self.save_dir)
util.save_object(self, stats_path)
util.save_checkpoint(self.save_dir,
iteration,
generative_model=generative_model,
inference_network=inference_network)
if iteration % self.eval_interval == 0:
log_p, kl = eval_gen_inf(generative_model, inference_network,
self.test_data_loader,
self.eval_num_particles)
_, renyi = eval_gen_inf_alpha(generative_model, inference_network,
self.test_data_loader,
self.eval_num_particles, self.alpha)
self.log_p_history.append(log_p)
self.kl_history.append(kl)
self.renyi_history.append(renyi)
stats = util.OnlineMeanStd()
for _ in range(10):
generative_model.zero_grad()
inference_network.zero_grad()
loss, elbo = losses.get_thermo_alpha_loss(
generative_model, inference_network, self.test_obs,
self.partition, self.num_particles, self.alpha,
self.integration)
loss.backward()
stats.update([p.grad for p in generative_model.parameters()] +
[p.grad for p in inference_network.parameters()])
self.grad_std_history.append(stats.avg_of_means_stds()[1].item())
util.print_with_time(
'Iteration {} log_p = {:.3f}, kl = {:.3f}, renyi = {:.3f}'.
format(iteration, self.log_p_history[-1], self.kl_history[-1],
self.renyi_history[-1]))
def process_feed(self, feed):
for entity in feed:
vehicle, trip_id = entity.vehicle, entity.vehicle.trip.trip_id
if vehicle.current_status == gtfs_realtime_pb2.VehiclePosition.STOPPED_AT:
if vehicle.current_stop_sequence == R_77_N_STOP_NUMBER:
if trip_id not in self.trains: # resilient against repeat feeds / long dwell times
self.trains[trip_id] = datetime.fromtimestamp(
entity.vehicle.timestamp)
elif vehicle.current_stop_sequence == R_36_N_STOP_NUMBER:
if trip_id in self.trains:
arrival_time = datetime.fromtimestamp(
entity.vehicle.timestamp)
departure_time = self.trains[trip_id]
print_with_time(
"Train {}: 77 at {}, 36 at {}, trip time {:.2f} minutes"
.format(trip_id, departure_time, arrival_time,
(arrival_time - departure_time).seconds /
60))
del self.trains[trip_id]
def find_arrived_trains_in_feed(self, feed):
arrived_trains_in_feed = {}
for entity in feed:
vehicle, trip = entity.vehicle, entity.vehicle.trip
if vehicle.current_stop_sequence == R_77_N_STOP_NUMBER and \
vehicle.current_status == gtfs_realtime_pb2.VehiclePosition.STOPPED_AT:
if trip.trip_id not in self.arrived_trains: # resilient against repeat feeds / long dwell times
arrived_trains_in_feed[
trip.trip_id] = datetime.fromtimestamp(
entity.vehicle.timestamp)
elif vehicle.current_stop_sequence > R_77_N_STOP_NUMBER and \
trip.trip_id not in self.arrived_trains and \
trip.trip_id not in self.alerted_trains:
print_with_time(
'Train {} currently at stop {} was never marked as arrived'
.format(trip.trip_id, vehicle.current_stop_sequence))
self.alerted_trains.add(trip.trip_id)
# todo handle short dwell times (i.e., no feed where STOPPED_AT R_77_N_STOP_NUMBER) using former approach
self.arrived_trains.update(arrived_trains_in_feed)
return arrived_trains_in_feed
def __call__(self, iteration, loss, elbo, generative_model,
inference_network, optimizer):
if iteration % self.logging_interval == 0:
util.print_with_time(
'Iteration {} loss = {:.3f}, elbo = {:.3f}'.format(
iteration, loss, elbo))
self.loss_history.append(loss)
self.elbo_history.append(elbo)
if iteration % self.checkpoint_interval == 0:
stats_filename = util.get_stats_path(self.model_folder)
util.save_object(self, stats_filename)
util.save_models(generative_model, inference_network,
self.model_folder, iteration)
if iteration % self.eval_interval == 0:
self.p_error_history.append(
util.get_p_error(self.true_generative_model, generative_model))
self.q_error_history.append(
util.get_q_error(self.true_generative_model, inference_network,
self.test_obss))
stats = util.OnlineMeanStd()
for _ in range(10):
inference_network.zero_grad()
if self.train_mode == 'vimco':
loss, elbo = losses.get_vimco_loss(generative_model,
inference_network,
self.test_obss,
self.num_particles)
elif self.train_mode == 'reinforce':
loss, elbo = losses.get_reinforce_loss(
generative_model, inference_network, self.test_obss,
self.num_particles)
loss.backward()
stats.update([p.grad for p in inference_network.parameters()])
self.grad_std_history.append(stats.avg_of_means_stds()[1])
util.print_with_time(
'Iteration {} p_error = {:.3f}, q_error_to_true = '
'{:.3f}'.format(iteration, self.p_error_history[-1],
self.q_error_history[-1]))
def start(self):
with open('api.key', 'r') as f:
api_key = f.read()
while True:
try:
feed = download_feed(
api_key,
self.feed_filter.lines[0]) # todo handle multiple lines
feed = self.feed_filter.filter(feed)
self.feed_processor.process_feed(feed)
except DecodeError as e:
print_with_time('Decode error!', e)
# todo print what the actual problem was
except TimeoutError as e:
print_with_time('Timeout error!', e)
except URLError as e:
print_with_time('URLError error!', e)
time.sleep(self.poll_frequency)
def print_arrived_trains(self, trains):
for key, value in sorted(
trains.items(),
key=lambda kv: kv[1]): # order by arrival time ascend
if len(self.arrived_trains
) == 1: # don't have headway for first train
print_with_time(
'Northbound R train {} arrived at 77th St station at {}'.
format(key, value))
else:
previous_arrival = self.arrived_trains[list(
self.arrived_trains)[-2]]
if value < previous_arrival: # todo better handling for order mix-ups
print_with_time(
'Out of order! Train {} at {} before most recent at {}'
.format(key, value, previous_arrival))
else:
print_with_time(
'Northbound R train {} arrived at 77th St station at {}, {:.2f} minutes after '
'most recent train.'.format(
key, value,
(value - previous_arrival).seconds / 60))
def run(args):
# set up args
if args.cuda and torch.cuda.is_available():
device = torch.device('cuda')
args.cuda = True
else:
device = torch.device('cpu')
args.cuda = False
if args.train_mode == 'thermo' or args.train_mode == 'thermo_wake':
partition = util.get_partition(args.num_partitions,
args.partition_type, args.log_beta_min,
device)
util.print_with_time('device = {}'.format(device))
util.print_with_time(str(args))
# save args
save_dir = util.get_save_dir()
args_path = util.get_args_path(save_dir)
util.save_object(args, args_path)
# data
binarized_mnist_train, binarized_mnist_valid, binarized_mnist_test = \
data.load_binarized_mnist(where=args.where)
data_loader = data.get_data_loader(binarized_mnist_train, args.batch_size,
device)
valid_data_loader = data.get_data_loader(binarized_mnist_valid,
args.valid_batch_size, device)
test_data_loader = data.get_data_loader(binarized_mnist_test,
args.test_batch_size, device)
train_obs_mean = torch.tensor(np.mean(binarized_mnist_train, axis=0),
device=device,
dtype=torch.float)
# init models
util.set_seed(args.seed)
generative_model, inference_network = util.init_models(
train_obs_mean, args.architecture, device)
# optim
optim_kwargs = {'lr': args.learning_rate}
# train
if args.train_mode == 'ws':
train_callback = train.TrainWakeSleepCallback(
save_dir, args.num_particles * args.batch_size, test_data_loader,
args.eval_num_particles, args.logging_interval,
args.checkpoint_interval, args.eval_interval)
train.train_wake_sleep(generative_model, inference_network,
data_loader, args.num_iterations,
args.num_particles, optim_kwargs,
train_callback)
elif args.train_mode == 'ww':
train_callback = train.TrainWakeWakeCallback(
save_dir, args.num_particles, test_data_loader,
args.eval_num_particles, args.logging_interval,
args.checkpoint_interval, args.eval_interval)
train.train_wake_wake(generative_model, inference_network, data_loader,
args.num_iterations, args.num_particles,
optim_kwargs, train_callback)
elif args.train_mode == 'reinforce' or args.train_mode == 'vimco':
train_callback = train.TrainIwaeCallback(
save_dir, args.num_particles, args.train_mode, test_data_loader,
args.eval_num_particles, args.logging_interval,
args.checkpoint_interval, args.eval_interval)
train.train_iwae(args.train_mode, generative_model, inference_network,
data_loader, args.num_iterations, args.num_particles,
optim_kwargs, train_callback)
elif args.train_mode == 'thermo':
train_callback = train.TrainThermoCallback(
save_dir, args.num_particles, partition, test_data_loader,
args.eval_num_particles, args.logging_interval,
args.checkpoint_interval, args.eval_interval)
train.train_thermo(generative_model, inference_network, data_loader,
args.num_iterations, args.num_particles, partition,
optim_kwargs, train_callback)
elif args.train_mode == 'thermo_wake':
train_callback = train.TrainThermoWakeCallback(
save_dir, args.num_particles, test_data_loader,
args.eval_num_particles, args.logging_interval,
args.checkpoint_interval, args.eval_interval)
train.train_thermo_wake(generative_model, inference_network,
data_loader, args.num_iterations,
args.num_particles, partition, optim_kwargs,
train_callback)
# eval validation
train_callback.valid_log_p, train_callback.valid_kl = train.eval_gen_inf(
generative_model, inference_network, valid_data_loader,
args.eval_num_particles)
# save models and stats
util.save_checkpoint(save_dir,
iteration=None,
generative_model=generative_model,
inference_network=inference_network)
stats_path = util.get_stats_path(save_dir)
util.save_object(train_callback, stats_path)
def get_mean_stds(generative_model, inference_network, num_mc_samples, obss,
num_particles):
vimco_grad = util.OnlineMeanStd()
vimco_one_grad = util.OnlineMeanStd()
reinforce_grad = util.OnlineMeanStd()
reinforce_one_grad = util.OnlineMeanStd()
two_grad = util.OnlineMeanStd()
log_evidence_stats = util.OnlineMeanStd()
log_evidence_grad = util.OnlineMeanStd()
wake_phi_loss_grad = util.OnlineMeanStd()
log_Q_grad = util.OnlineMeanStd()
sleep_loss_grad = util.OnlineMeanStd()
for mc_sample_idx in range(num_mc_samples):
util.print_with_time('MC sample {}'.format(mc_sample_idx))
log_weight, log_q = losses.get_log_weight_and_log_q(
generative_model, inference_network, obss, num_particles)
log_evidence = torch.logsumexp(log_weight, dim=1) - \
np.log(num_particles)
avg_log_evidence = torch.mean(log_evidence)
log_Q = torch.sum(log_q, dim=1)
avg_log_Q = torch.mean(log_Q)
reinforce_one = torch.mean(log_evidence.detach() * log_Q)
reinforce = reinforce_one + avg_log_evidence
vimco_one = 0
for i in range(num_particles):
log_weight_ = log_weight[:, util.range_except(num_particles, i)]
control_variate = torch.logsumexp(
torch.cat([log_weight_, torch.mean(log_weight_, dim=1,
keepdim=True)], dim=1),
dim=1)
vimco_one = vimco_one + (log_evidence.detach() -
control_variate.detach()) * log_q[:, i]
vimco_one = torch.mean(vimco_one)
vimco = vimco_one + avg_log_evidence
normalized_weight = util.exponentiate_and_normalize(log_weight, dim=1)
wake_phi_loss = torch.mean(
-torch.sum(normalized_weight.detach() * log_q, dim=1))
inference_network.zero_grad()
generative_model.zero_grad()
vimco.backward(retain_graph=True)
vimco_grad.update([param.grad for param in
inference_network.parameters()])
inference_network.zero_grad()
generative_model.zero_grad()
vimco_one.backward(retain_graph=True)
vimco_one_grad.update([param.grad for param in
inference_network.parameters()])
inference_network.zero_grad()
generative_model.zero_grad()
reinforce.backward(retain_graph=True)
reinforce_grad.update([param.grad for param in
inference_network.parameters()])
inference_network.zero_grad()
generative_model.zero_grad()
reinforce_one.backward(retain_graph=True)
reinforce_one_grad.update([param.grad for param in
inference_network.parameters()])
inference_network.zero_grad()
generative_model.zero_grad()
avg_log_evidence.backward(retain_graph=True)
two_grad.update([param.grad for param in
inference_network.parameters()])
log_evidence_grad.update([param.grad for param in
generative_model.parameters()])
inference_network.zero_grad()
generative_model.zero_grad()
wake_phi_loss.backward(retain_graph=True)
wake_phi_loss_grad.update([param.grad for param in
inference_network.parameters()])
inference_network.zero_grad()
generative_model.zero_grad()
avg_log_Q.backward(retain_graph=True)
log_Q_grad.update([param.grad for param in
inference_network.parameters()])
log_evidence_stats.update([avg_log_evidence.unsqueeze(0)])
sleep_loss = losses.get_sleep_loss(
generative_model, inference_network, num_particles * len(obss))
inference_network.zero_grad()
generative_model.zero_grad()
sleep_loss.backward()
sleep_loss_grad.update([param.grad for param in
inference_network.parameters()])
return list(map(
lambda x: x.avg_of_means_stds(),
[vimco_grad, vimco_one_grad, reinforce_grad, reinforce_one_grad,
two_grad, log_evidence_stats, log_evidence_grad, wake_phi_loss_grad,
log_Q_grad, sleep_loss_grad]))
def __call__(self, iteration, sleep_loss, generative_model,
inference_network, optimizer):
if iteration % self.logging_interval == 0:
util.print_with_time('Iteration {}: loss = {:.3f}'.format(
iteration, sleep_loss))
self.sleep_loss_history.append(sleep_loss)
def train_vimco(
generative_model,
inference_network,
data_loader,
num_iterations,
num_particles,
true_cluster_cov,
test_data_loader,
test_num_particles,
true_generative_model,
checkpoint_path,
):
optimizer = torch.optim.Adam(
itertools.chain(generative_model.parameters(),
inference_network.parameters()))
(
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
) = ([], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [])
data_loader_iter = iter(data_loader)
for iteration in range(num_iterations):
# get obs
try:
obs = next(data_loader_iter)
except StopIteration:
data_loader_iter = iter(data_loader)
obs = next(data_loader_iter)
# loss
optimizer.zero_grad()
loss, elbo = losses.get_vimco_loss(generative_model, inference_network,
obs, num_particles)
loss.backward(retain_graph=True)
optimizer.step()
theta_losses.append(loss.item())
phi_losses.append(loss.item())
cluster_cov_distances.append(
torch.norm(true_cluster_cov -
generative_model.get_cluster_cov()).item())
if iteration % 100 == 0: # test every 100 iterations
(
test_log_p,
test_log_p_true,
test_kl_qp,
test_kl_pq,
test_kl_qp_true,
test_kl_pq_true,
_,
_,
) = models.eval_gen_inf(true_generative_model, generative_model,
inference_network, None, test_data_loader)
test_log_ps.append(test_log_p)
test_log_ps_true.append(test_log_p_true)
test_kl_qps.append(test_kl_qp)
test_kl_pqs.append(test_kl_pq)
test_kl_qps_true.append(test_kl_qp_true)
test_kl_pqs_true.append(test_kl_pq_true)
(
train_log_p,
train_log_p_true,
train_kl_qp,
train_kl_pq,
train_kl_qp_true,
train_kl_pq_true,
_,
_,
reweighted_train_kl_qp,
reweighted_train_kl_qp_true,
) = models.eval_gen_inf(
true_generative_model,
generative_model,
inference_network,
None,
data_loader,
num_particles=num_particles,
reweighted_kl=True,
)
train_log_ps.append(train_log_p)
train_log_ps_true.append(train_log_p_true)
train_kl_qps.append(train_kl_qp)
train_kl_pqs.append(train_kl_pq)
train_kl_qps_true.append(train_kl_qp_true)
train_kl_pqs_true.append(train_kl_pq_true)
reweighted_train_kl_qps.append(reweighted_train_kl_qp)
reweighted_train_kl_qps_true.append(reweighted_train_kl_qp_true)
util.save_checkpoint(
checkpoint_path,
generative_model,
inference_network,
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
None,
None,
None,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
)
util.print_with_time("it. {} | theta loss = {:.2f}".format(
iteration, loss))
# if iteration % 200 == 0:
# z = inference_network.get_latent_dist(obs).sample()
# util.save_plot("images/rws/iteration_{}.png".format(iteration),
# obs[:3], z[:3])
return (
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
None,
None,
None,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
)
if __name__ == "__main__":
import os
try:
os.remove("example_learner_data.json")
except:
pass
try:
os.remove("./saved_estimators/fluent_afm.pkl")
except:
pass
try:
os.remove("./saved_estimators/fluent_online_transfer_learning.pkl")
except:
pass
with open('xapi_statements_fluent.json', 'r') as infile:
xapi_statements = json.load(infile)
# TODO: Find out why processing multiple statements at once work, but not code below that runs one at a time (AFM error).
# print(train_and_predict('fluent', xapi_statements))
# print("---")
# statements = [{"version": "1.0.0", "id": "7df5a88b-a527-11e7-93e8-9cb6d063d611", "actor": {"objectType": "Agent", "name": "D80FD", "account": {"name": "D80FD", "homePage": "https://pslcdatashop.web.cmu.edu/DatasetInfo?datasetId=263"}}, "verb": {"id": "http://adlnet.gov/expapi/verbs/passed", "display": {"en-US": "passed"}}, "object": {"objectType": "Activity", "id": "https://pslcdatashop.web.cmu.edu/DatasetInfo?datasetId=263/step/%28BODY%20SWIMMER%29"}, "context": {"contextActivities": {"parent": [{"objectType": "Activity", "id": "https://pslcdatashop.web.cmu.edu/DatasetInfo?datasetId=263/problem/VEC1A"}], "grouping": [{"objectType": "Activity", "id": "https://pslcdatashop.web.cmu.edu/DatasetInfo?datasetId=263/problem/VEC1A"}]}, "extensions": {"http://www.soartech.com/target/learnsphere-transaction-id": "436fd8ae3a41f4646677b3071f23d37a", "http://www.soartech.com/target/competencies": [{"competencyId": "http://insertCassUrl/api/custom/data/insertCassSchemaUrl.0.2.Competency/pslcdatashop.web.cmu.edu/DatasetInfo?datasetId=263/Default/DRAW-BODY", "framework": "Default"}]}}, "result": {"duration": "PT2.4S", "success": True}, "timestamp": "2008-09-10T00:50:52+00:00", "authority": {"objectType": "Agent", "name": "API Client", "mbox": "mailto:[email protected]"}, "stored": "2017-09-29T15:40:50.986800+00:00"}]
for statement in xapi_statements:
print_with_time(train_and_predict('fluent', [statement]))
print("---")
#run_afm(statements)
def run(args):
# set up args
args.device = None
if args.cuda and torch.cuda.is_available():
args.device = torch.device('cuda')
else:
args.device = torch.device('cpu')
args.num_mixtures = 20
if args.init_near:
args.init_mixture_logits = np.ones(args.num_mixtures)
else:
args.init_mixture_logits = np.array(
list(reversed(2 * np.arange(args.num_mixtures))))
args.softmax_multiplier = 0.5
if args.train_mode == 'concrete':
args.relaxed_one_hot = True
args.temperature = 3
else:
args.relaxed_one_hot = False
args.temperature = None
temp = np.arange(args.num_mixtures) + 5
true_p_mixture_probs = temp / np.sum(temp)
args.true_mixture_logits = \
np.log(true_p_mixture_probs) / args.softmax_multiplier
util.print_with_time(str(args))
# save args
model_folder = util.get_model_folder()
args_filename = util.get_args_path(model_folder)
util.save_object(args, args_filename)
# init models
util.set_seed(args.seed)
generative_model, inference_network, true_generative_model = \
util.init_models(args)
if args.train_mode == 'relax':
control_variate = models.ControlVariate(args.num_mixtures)
# init dataloader
obss_data_loader = torch.utils.data.DataLoader(
true_generative_model.sample_obs(args.num_obss),
batch_size=args.batch_size,
shuffle=True)
# train
if args.train_mode == 'mws':
train_callback = train.TrainMWSCallback(model_folder,
true_generative_model,
args.logging_interval,
args.checkpoint_interval,
args.eval_interval)
train.train_mws(generative_model, inference_network, obss_data_loader,
args.num_iterations, args.mws_memory_size,
train_callback)
if args.train_mode == 'ws':
train_callback = train.TrainWakeSleepCallback(
model_folder, true_generative_model,
args.batch_size * args.num_particles, args.logging_interval,
args.checkpoint_interval, args.eval_interval)
train.train_wake_sleep(generative_model, inference_network,
obss_data_loader, args.num_iterations,
args.num_particles, train_callback)
elif args.train_mode == 'ww':
train_callback = train.TrainWakeWakeCallback(model_folder,
true_generative_model,
args.num_particles,
args.logging_interval,
args.checkpoint_interval,
args.eval_interval)
train.train_wake_wake(generative_model, inference_network,
obss_data_loader, args.num_iterations,
args.num_particles, train_callback)
elif args.train_mode == 'dww':
train_callback = train.TrainDefensiveWakeWakeCallback(
model_folder, true_generative_model, args.num_particles, 0.2,
args.logging_interval, args.checkpoint_interval,
args.eval_interval)
train.train_defensive_wake_wake(0.2, generative_model,
inference_network, obss_data_loader,
args.num_iterations,
args.num_particles, train_callback)
elif args.train_mode == 'reinforce' or args.train_mode == 'vimco':
train_callback = train.TrainIwaeCallback(
model_folder, true_generative_model, args.num_particles,
args.train_mode, args.logging_interval, args.checkpoint_interval,
args.eval_interval)
train.train_iwae(args.train_mode, generative_model, inference_network,
obss_data_loader, args.num_iterations,
args.num_particles, train_callback)
elif args.train_mode == 'concrete':
train_callback = train.TrainConcreteCallback(
model_folder, true_generative_model, args.num_particles,
args.num_iterations, args.logging_interval,
args.checkpoint_interval, args.eval_interval)
train.train_iwae(args.train_mode, generative_model, inference_network,
obss_data_loader, args.num_iterations,
args.num_particles, train_callback)
elif args.train_mode == 'relax':
train_callback = train.TrainRelaxCallback(model_folder,
true_generative_model,
args.num_particles,
args.logging_interval,
args.checkpoint_interval,
args.eval_interval)
train.train_relax(generative_model, inference_network, control_variate,
obss_data_loader, args.num_iterations,
args.num_particles, train_callback)
# save models and stats
util.save_models(generative_model, inference_network, model_folder)
if args.train_mode == 'relax':
util.save_control_variate(control_variate, model_folder)
stats_filename = util.get_stats_path(model_folder)
util.save_object(train_callback, stats_filename)
def train_mws(
generative_model,
inference_network,
data_loader,
num_iterations,
memory_size,
true_cluster_cov,
test_data_loader,
test_num_particles,
true_generative_model,
checkpoint_path,
reweighted=False,
):
optimizer = torch.optim.Adam(
itertools.chain(generative_model.parameters(),
inference_network.parameters()))
(
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
) = ([], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [])
memory = {}
data_loader_iter = iter(data_loader)
for iteration in range(num_iterations):
# get obs
try:
obs = next(data_loader_iter)
except StopIteration:
data_loader_iter = iter(data_loader)
obs = next(data_loader_iter)
theta_loss = 0
phi_loss = 0
for single_obs in obs:
# key to index memory
single_obs_key = tuple(single_obs.tolist())
# populate memory if empty
if (single_obs_key not in memory) or len(
memory[single_obs_key]) == 0:
# batch shape [1] and event shape [num_data]
latent_dist = inference_network.get_latent_dist(
single_obs.unsqueeze(0))
# HACK
while True:
# [memory_size, num_data]
latent = inference_network.sample_from_latent_dist(
latent_dist, memory_size).squeeze(1)
# list of M \in {1, ..., memory_size} elements
# could be less than memory_size because
# sampled elements can be duplicate
memory[single_obs_key] = list(
set([tuple(x.tolist()) for x in latent]))
if len(memory[single_obs_key]) == memory_size:
break
# WAKE
# batch shape [1] and event shape [num_data]
latent_dist = inference_network.get_latent_dist(
single_obs.unsqueeze(0))
# [1, 1, num_data] -> [num_data]
latent = inference_network.sample_from_latent_dist(latent_dist,
1).view(-1)
# set (of size memory_size + 1) of tuples (of length num_data)
memoized_latent_plus_current_latent = set(
memory.get(single_obs_key, []) + [tuple(latent.tolist())])
# [memory_size + 1, 1, num_data]
memoized_latent_plus_current_latent_tensor = torch.tensor(
list(memoized_latent_plus_current_latent),
device=single_obs.device).unsqueeze(1)
# [memory_size + 1]
log_p_tensor = generative_model.get_log_prob(
memoized_latent_plus_current_latent_tensor,
single_obs.unsqueeze(0)).squeeze(-1)
# this takes the longest
# {int: [], ...}
log_p = {
mem_latent: lp
for mem_latent, lp in zip(memoized_latent_plus_current_latent,
log_p_tensor)
}
# update memory.
# {float: list of ints}
memory[single_obs_key] = sorted(
memoized_latent_plus_current_latent,
key=log_p.get)[-memory_size:]
# REMEMBER
# []
if reweighted:
memory_log_weight = torch.stack(
list(map(log_p.get,
memory[single_obs_key]))) # [memory_size]
memory_weight_normalized = util.exponentiate_and_normalize(
memory_log_weight, dim=0) # [memory_size]
memory_latent = torch.tensor(
memory[single_obs_key]) # [memory_size, num_data]
inference_network_log_prob = inference_network.get_log_prob_from_latent_dist(
latent_dist,
memory_latent[:, None, :]).squeeze(-1) # [memory_size]
theta_loss += -torch.sum(
memory_log_weight *
memory_weight_normalized.detach()) / len(obs)
phi_loss += -torch.sum(
inference_network_log_prob *
memory_weight_normalized.detach()) / len(obs)
else:
remembered_latent_id_dist = torch.distributions.Categorical(
logits=torch.tensor(
list(map(log_p.get, memory[single_obs_key]))))
remembered_latent_id = remembered_latent_id_dist.sample()
remembered_latent_id_log_prob = remembered_latent_id_dist.log_prob(
remembered_latent_id)
remembered_latent = memory[single_obs_key][
remembered_latent_id]
remembered_latent_tensor = torch.tensor(
[remembered_latent], device=single_obs.device)
# []
theta_loss += -(log_p.get(remembered_latent) -
remembered_latent_id_log_prob.detach()) / len(
obs)
# []
phi_loss += -inference_network.get_log_prob_from_latent_dist(
latent_dist, remembered_latent_tensor).view(()) / len(obs)
# SLEEP
# TODO
optimizer.zero_grad()
theta_loss.backward()
phi_loss.backward()
optimizer.step()
theta_losses.append(theta_loss.item())
phi_losses.append(phi_loss.item())
cluster_cov_distances.append(
torch.norm(true_cluster_cov -
generative_model.get_cluster_cov()).item())
if iteration % 100 == 0: # test every 100 iterations
(
test_log_p,
test_log_p_true,
test_kl_qp,
test_kl_pq,
test_kl_qp_true,
test_kl_pq_true,
_,
_,
) = models.eval_gen_inf(true_generative_model, generative_model,
inference_network, None, test_data_loader)
test_log_ps.append(test_log_p)
test_log_ps_true.append(test_log_p_true)
test_kl_qps.append(test_kl_qp)
test_kl_pqs.append(test_kl_pq)
test_kl_qps_true.append(test_kl_qp_true)
test_kl_pqs_true.append(test_kl_pq_true)
(
train_log_p,
train_log_p_true,
train_kl_qp,
train_kl_pq,
train_kl_qp_true,
train_kl_pq_true,
train_kl_memory_p,
train_kl_memory_p_true,
) = models.eval_gen_inf(true_generative_model, generative_model,
inference_network, memory, data_loader)
train_log_ps.append(train_log_p)
train_log_ps_true.append(train_log_p_true)
train_kl_qps.append(train_kl_qp)
train_kl_pqs.append(train_kl_pq)
train_kl_qps_true.append(train_kl_qp_true)
train_kl_pqs_true.append(train_kl_pq_true)
train_kl_memory_ps.append(train_kl_memory_p)
train_kl_memory_ps_true.append(train_kl_memory_p_true)
util.save_checkpoint(
checkpoint_path,
generative_model,
inference_network,
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
memory,
None,
None,
)
util.print_with_time(
"it. {} | theta loss = {:.2f} | phi loss = {:.2f}".format(
iteration, theta_loss, phi_loss))
# if iteration % 200 == 0:
# z = inference_network.get_latent_dist(obs).sample()
# util.save_plot("images/mws/iteration_{}.png".format(iteration),
# obs[:3], z[:3])
return (
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
memory,
None,
None,
)
def run(args):
# set up args
if args.cuda and torch.cuda.is_available():
device = torch.device("cuda")
args.cuda = True
else:
device = torch.device("cpu")
args.cuda = False
util.print_with_time("args = {}".format(args))
# init
util.print_with_time("init")
true_cluster_cov = torch.eye(args.num_dim, device=device) * 0.03
generative_model, inference_network, true_generative_model = util.init(
args.num_data, args.num_dim, true_cluster_cov, device
)
util.set_seed(args.seed)
# data
util.print_with_time("data")
train_data, test_data = util.load_data()
data_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=True)
test_data_loader = torch.utils.data.DataLoader(
test_data, batch_size=args.batch_size, shuffle=True
)
# train
util.print_with_time("train")
checkpoint_path = "checkpoints/{}_{}_{}_{}.pt".format(
args.checkpoint_path_prefix, args.algorithm, args.num_particles, args.seed
)
Path("checkpoints/").mkdir(parents=True, exist_ok=True)
if args.algorithm == "mws":
(
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
) = train.train_mws(
generative_model,
inference_network,
data_loader,
args.num_iterations,
args.memory_size,
true_cluster_cov,
test_data_loader,
args.test_num_particles,
true_generative_model,
checkpoint_path,
)
elif args.algorithm == "rmws":
(
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
) = train.train_mws(
generative_model,
inference_network,
data_loader,
args.num_iterations,
args.memory_size,
true_cluster_cov,
test_data_loader,
args.test_num_particles,
true_generative_model,
checkpoint_path,
reweighted=True,
)
elif args.algorithm == "rws":
(
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
) = train.train_rws(
generative_model,
inference_network,
data_loader,
args.num_iterations,
args.num_particles,
true_cluster_cov,
test_data_loader,
args.test_num_particles,
true_generative_model,
checkpoint_path,
)
elif args.algorithm == "vimco":
(
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
) = train.train_vimco(
generative_model,
inference_network,
data_loader,
args.num_iterations,
args.num_particles,
true_cluster_cov,
test_data_loader,
args.test_num_particles,
true_generative_model,
checkpoint_path,
)
# save model
util.save_checkpoint(
checkpoint_path,
generative_model,
inference_network,
theta_losses,
phi_losses,
cluster_cov_distances,
test_log_ps,
test_log_ps_true,
test_kl_qps,
test_kl_pqs,
test_kl_qps_true,
test_kl_pqs_true,
train_log_ps,
train_log_ps_true,
train_kl_qps,
train_kl_pqs,
train_kl_qps_true,
train_kl_pqs_true,
train_kl_memory_ps,
train_kl_memory_ps_true,
memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true,
)
def run(args):
# set up args
if args.cuda and torch.cuda.is_available():
device = torch.device('cuda')
args.cuda = True
else:
device = torch.device('cpu')
args.cuda = False
util.print_with_time('args = {}'.format(args))
# init
util.print_with_time('init')
true_cluster_cov = torch.eye(args.num_dim, device=device) * 0.03
generative_model, inference_network, true_generative_model = util.init(
args.num_data, args.num_dim, true_cluster_cov, device)
util.set_seed(args.seed)
# data
util.print_with_time('data')
train_data, test_data = util.load_data()
data_loader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True)
test_data_loader = torch.utils.data.DataLoader(test_data,
batch_size=args.batch_size,
shuffle=True)
# train
util.print_with_time('train')
if args.test_run:
checkpoint_path = 'checkpoint.pt'
else:
checkpoint_path = 'checkpoints/{}_{}_{}_{}.pt'.format(
args.checkpoint_path_prefix, args.algorithm, args.seed,
args.num_particles)
if args.algorithm == 'mws':
(theta_losses, phi_losses, cluster_cov_distances, test_log_ps,
test_log_ps_true, test_kl_qps, test_kl_pqs, test_kl_qps_true,
test_kl_pqs_true, train_log_ps, train_log_ps_true, train_kl_qps,
train_kl_pqs, train_kl_qps_true, train_kl_pqs_true,
train_kl_memory_ps, train_kl_memory_ps_true, memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true) = train.train_mws(
generative_model, inference_network, data_loader,
args.num_iterations, args.memory_size, true_cluster_cov,
test_data_loader, args.test_num_particles, true_generative_model,
checkpoint_path)
elif args.algorithm == 'rmws':
(theta_losses, phi_losses, cluster_cov_distances, test_log_ps,
test_log_ps_true, test_kl_qps, test_kl_pqs, test_kl_qps_true,
test_kl_pqs_true, train_log_ps, train_log_ps_true, train_kl_qps,
train_kl_pqs, train_kl_qps_true, train_kl_pqs_true,
train_kl_memory_ps, train_kl_memory_ps_true, memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true) = train.train_mws(
generative_model,
inference_network,
data_loader,
args.num_iterations,
args.memory_size,
true_cluster_cov,
test_data_loader,
args.test_num_particles,
true_generative_model,
checkpoint_path,
reweighted=True)
elif args.algorithm == 'rws':
(theta_losses, phi_losses, cluster_cov_distances, test_log_ps,
test_log_ps_true, test_kl_qps, test_kl_pqs, test_kl_qps_true,
test_kl_pqs_true, train_log_ps, train_log_ps_true, train_kl_qps,
train_kl_pqs, train_kl_qps_true, train_kl_pqs_true,
train_kl_memory_ps, train_kl_memory_ps_true, memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true) = train.train_rws(
generative_model, inference_network, data_loader,
args.num_iterations, args.num_particles, true_cluster_cov,
test_data_loader, args.test_num_particles, true_generative_model,
checkpoint_path)
elif args.algorithm == 'vimco':
(theta_losses, phi_losses, cluster_cov_distances, test_log_ps,
test_log_ps_true, test_kl_qps, test_kl_pqs, test_kl_qps_true,
test_kl_pqs_true, train_log_ps, train_log_ps_true, train_kl_qps,
train_kl_pqs, train_kl_qps_true, train_kl_pqs_true,
train_kl_memory_ps, train_kl_memory_ps_true, memory,
reweighted_train_kl_qps,
reweighted_train_kl_qps_true) = train.train_vimco(
generative_model, inference_network, data_loader,
args.num_iterations, args.num_particles, true_cluster_cov,
test_data_loader, args.test_num_particles, true_generative_model,
checkpoint_path)
# save model
util.save_checkpoint(checkpoint_path, generative_model, inference_network,
theta_losses, phi_losses, cluster_cov_distances,
test_log_ps, test_log_ps_true, test_kl_qps,
test_kl_pqs, test_kl_qps_true, test_kl_pqs_true,
train_log_ps, train_log_ps_true, train_kl_qps,
train_kl_pqs, train_kl_qps_true, train_kl_pqs_true,
train_kl_memory_ps, train_kl_memory_ps_true, memory,
reweighted_train_kl_qps, reweighted_train_kl_qps_true)
def train_rws(generative_model, inference_network, data_loader, num_iterations,
num_particles, true_cluster_cov, test_data_loader,
test_num_particles, true_generative_model, checkpoint_path):
optimizer_phi = torch.optim.Adam(inference_network.parameters())
optimizer_theta = torch.optim.Adam(generative_model.parameters())
(theta_losses, phi_losses, cluster_cov_distances,
test_log_ps, test_log_ps_true, test_kl_qps, test_kl_pqs, test_kl_qps_true, test_kl_pqs_true,
train_log_ps, train_log_ps_true, train_kl_qps, train_kl_pqs, train_kl_qps_true, train_kl_pqs_true,
reweighted_train_kl_qps, reweighted_train_kl_qps_true) = \
[], [], [], [], [], [], [], [], [], [], [], [], [], [], [], [], []
data_loader_iter = iter(data_loader)
for iteration in range(num_iterations):
# get obs
try:
obs = next(data_loader_iter)
except StopIteration:
data_loader_iter = iter(data_loader)
obs = next(data_loader_iter)
log_weight, log_q = losses.get_log_weight_and_log_q(
generative_model, inference_network, obs, num_particles)
# wake theta
optimizer_phi.zero_grad()
optimizer_theta.zero_grad()
wake_theta_loss, elbo = losses.get_wake_theta_loss_from_log_weight(
log_weight)
wake_theta_loss.backward(retain_graph=True)
optimizer_theta.step()
# wake phi
optimizer_phi.zero_grad()
optimizer_theta.zero_grad()
wake_phi_loss = losses.get_wake_phi_loss_from_log_weight_and_log_q(
log_weight, log_q)
wake_phi_loss.backward()
optimizer_phi.step()
theta_losses.append(wake_theta_loss.item())
phi_losses.append(wake_phi_loss.item())
cluster_cov_distances.append(
torch.norm(true_cluster_cov -
generative_model.get_cluster_cov()).item())
if iteration % 100 == 0: # test every 100 iterations
(test_log_p, test_log_p_true, test_kl_qp, test_kl_pq,
test_kl_qp_true, test_kl_pq_true, _,
_) = models.eval_gen_inf(true_generative_model, generative_model,
inference_network, None,
test_data_loader)
test_log_ps.append(test_log_p)
test_log_ps_true.append(test_log_p_true)
test_kl_qps.append(test_kl_qp)
test_kl_pqs.append(test_kl_pq)
test_kl_qps_true.append(test_kl_qp_true)
test_kl_pqs_true.append(test_kl_pq_true)
(train_log_p, train_log_p_true, train_kl_qp, train_kl_pq,
train_kl_qp_true, train_kl_pq_true, _, _, reweighted_train_kl_qp,
reweighted_train_kl_qp_true) = models.eval_gen_inf(
true_generative_model,
generative_model,
inference_network,
None,
data_loader,
num_particles=num_particles,
reweighted_kl=True)
train_log_ps.append(train_log_p)
train_log_ps_true.append(train_log_p_true)
train_kl_qps.append(train_kl_qp)
train_kl_pqs.append(train_kl_pq)
train_kl_qps_true.append(train_kl_qp_true)
train_kl_pqs_true.append(train_kl_pq_true)
reweighted_train_kl_qps.append(reweighted_train_kl_qp)
reweighted_train_kl_qps_true.append(reweighted_train_kl_qp_true)
util.save_checkpoint(
checkpoint_path, generative_model, inference_network,
theta_losses, phi_losses, cluster_cov_distances, test_log_ps,
test_log_ps_true, test_kl_qps, test_kl_pqs, test_kl_qps_true,
test_kl_pqs_true, train_log_ps, train_log_ps_true,
train_kl_qps, train_kl_pqs, train_kl_qps_true,
train_kl_pqs_true, None, None, None, reweighted_train_kl_qps,
reweighted_train_kl_qps_true)
util.print_with_time(
'it. {} | theta loss = {:.2f} | phi loss = {:.2f}'.format(
iteration, wake_theta_loss, wake_phi_loss))
# if iteration % 200 == 0:
# z = inference_network.get_latent_dist(obs).sample()
# util.save_plot("images/rws/iteration_{}.png".format(iteration),
# obs[:3], z[:3])
return (theta_losses, phi_losses, cluster_cov_distances, test_log_ps,
test_log_ps_true, test_kl_qps, test_kl_pqs, test_kl_qps_true,
test_kl_pqs_true, train_log_ps, train_log_ps_true, train_kl_qps,
train_kl_pqs, train_kl_qps_true, train_kl_pqs_true, None, None,
None, reweighted_train_kl_qps, reweighted_train_kl_qps_true)
