Manual LR scheduler and change printing.

# define stochastic loss function

loss_fn = lambda num_samples, manual_samples: -model.log_likelihood(X, Y, num_samples=num_samples, manual_samples=manual_samples)/X.shape[0]

# get optimizer and LR scheduler

# get optimizer

optimizer = optim.Adam(model.parameters(), lr=args.lr)

scheduler = optim.lr_scheduler.MultiplicativeLR(optimizer, lr_lambda=lambda iter: 0.9)

# initialize results dict

results_train_keys = ["v_b", "v_w", "v_n", "loss", "loss_new", "delta_loss"]

		delta_loss = loss_new-loss

	# decrease LR if loss increases

	if delta_loss.item() > 0:

		scheduler.step()

		for param_group in optimizer.param_groups:

			param_group["lr"] *= 0.9

	# update results

	for (key, value) in zip(results_train_keys, [model.v_b(), model.v_w(), model.v_n(), loss, loss_new, delta_loss]):

		results["train"][key].append( value.item() )

	# print loss info

	if i >= 5:

		print("iter {} delta loss {:.3f} % mov avg loss {:.3f}".format(i, 100*(loss_new/loss-1.0).item(), sum(results["train"]["loss"][-5:])/5))

	else:

		print("iter {} delta loss {:.3f} % mov avg loss ---".format(i, 100*(loss_new/loss-1.0).item()))

	print("iter {} loss {:.3f} delta loss {:.3f} %".format(i, loss.item(), 100*(loss_new/loss-1.0).item()))

# compute loss of final iteration

with torch.no_grad():

	loss = loss_fn(args.train_samples, None)

results["train"]["loss"].append(loss.item())

# record training time

# record and print training time

results["train"]["time"] = time.time()-time_0

print("Took {:.3f} seconds".format(results["train"]["time"]))

# get final values with higher number of samples

time_0 = time.time()