parser.add_argument("--grid_size", type=int, default=6)
+parser.add_argument("--grid_nb_colors", type=int, default=6)
+
+parser.add_argument("--grid_nb_shapes", type=int, default=6)
+
##############################
# picoclvr options
nb_test_samples=args.nb_test_samples,
batch_size=args.batch_size,
size=args.grid_size,
+ nb_shapes=args.grid_nb_shapes,
+ nb_colors=args.grid_nb_colors,
logger=log_string,
device=device_data,
)
##############################
-for input in task.batches(split="train", desc="calibrate"):
- input = input.to(device)
- output = model(mygpt.BracketedSequence(input)).x
-
-for n, m in model.named_modules():
- for a in dir(m):
- x = getattr(m, a)
- if isinstance(x, mygpt.Calibrator):
- print(f"####### ${n} | ${a} ########################")
- mean, std = x.moments()
- print("mean\n", mean, "\n")
- print("std\n", std, "\n")
- print(f"############################################\n\n")
-
-exit(0)
+if "calibrate" in sup_args:
+ for input in task.batches(split="train", desc="calibrate"):
+ input = input.to(device)
+ output = model(mygpt.BracketedSequence(input)).x
+
+ for n, m in model.named_modules():
+ for a in dir(m):
+ x = getattr(m, a)
+ if isinstance(x, mygpt.Calibrator):
+ print(f"####### ${n} | ${a} ########################")
+ mean, std = x.moments()
+ print("mean\n", mean, "\n")
+ print("std\n", std, "\n")
+ print(f"############################################\n\n")
+
+ exit(0)
##############################