Commit d96563cd authored by lucas_miranda's avatar lucas_miranda
Browse files

Added CustomStopper class to train_utils.py, to start early stopping only after annealing is over

parent ce193ba1
Pipeline #89043 passed with stage
in 21 minutes and 40 seconds
......@@ -122,6 +122,13 @@ parser.add_argument(
default=10,
type=int,
)
parser.add_argument(
"--montecarlo-kl",
"-mckl",
help="Number of samples to compute when adding KLDivergence to the loss function",
default=10,
type=int,
)
parser.add_argument(
"--neuron-control",
"-nc",
......@@ -219,6 +226,7 @@ k = args.components
kl_wu = args.kl_warmup
loss = args.loss
mmd_wu = args.mmd_warmup
mc_kl = args.montecarlo_kl
neuron_control = args.neuron_control
output_path = os.path.join(args.output_path)
overlap_loss = args.overlap_loss
......@@ -440,6 +448,7 @@ if not tune:
kl_warmup_epochs=kl_wu,
loss=loss,
mmd_warmup_epochs=mmd_wu,
montecarlo_kl=mc_kl,
neuron_control=neuron_control,
number_of_components=k,
overlap_loss=overlap_loss,
......
......@@ -76,6 +76,7 @@ rule explore_encoding_dimension_and_loss_function:
"--loss {wildcards.loss} "
"--kl-warmup 2 "
"--mmd-warmup 2 "
"--montecarlo-kl 10 "
"--encoding-size {wildcards.encs} "
"--batch-size 256 "
"--window-size 11 "
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment