########################################
-parser.add_argument("--nb_epochs", type=int, default=100)
+parser.add_argument("--nb_epochs", type=int, default=10000)
parser.add_argument("--batch_size", type=int, default=None)
parser.add_argument("--nb_test_samples", type=int, default=None)
-parser.add_argument("--optim", type=str, default="adam")
-
parser.add_argument("--learning_rate", type=float, default=1e-4)
-parser.add_argument("--learning_rate_schedule", type=str, default=None)
-
########################################
parser.add_argument("--model", type=str, default=None)
##############################
-if args.learning_rate_schedule == "cos":
- learning_rate_schedule = {}
- for n_epoch in range(args.nb_epochs):
- u = n_epoch / args.nb_epochs * math.pi
- learning_rate_schedule[n_epoch] = args.learning_rate * 0.5 * (1 + math.cos(u))
-else:
- if args.learning_rate_schedule is not None:
- u = {
- int(k): float(v)
- for k, v in [
- tuple(x.split(":")) for x in args.learning_rate_schedule.split(",")
- ]
- }
- else:
- u = {}
-
- learning_rate_schedule = {}
- learning_rate = args.learning_rate
- for n_epoch in range(args.nb_epochs):
- if n_epoch in u:
- learning_rate = u[n_epoch]
- learning_rate_schedule[n_epoch] = learning_rate
-
-log_string(f"learning_rate_schedule {learning_rate_schedule}")
-
-######################################################################
-
def one_epoch(model, task):
- if args.optim == "sgd":
- optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate)
- elif args.optim == "adam":
- optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
- elif args.optim == "adamw":
- optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
- else:
- raise ValueError(f"Unknown optimizer {args.optim}.")
+ optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
model.train()
task.save_image(
new_quizzes[:96],
args.result_dir,
- f"world_new_{n_epoch:04d}.png",
+ f"world_new_{n_epoch:04d}_{model.id:02d}.png",
log_string,
)