def main(config, comet=False):
config = Config(config)
# comet-ml setting
if comet:
experiment = Experiment(api_key=config.api_key,
project_name=config.project_name,
workspace=config.workspace)
experiment.log_parameters(config)
# device and dataset setting
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = MCDataset(config.root, config.dataset_name)
data = dataset[0].to(device)
config.num_nodes = dataset.num_nodes
config.num_users = int(data.num_users)
config.num_relations = dataset.num_relations # defines number of edge types
# set and init model
model = GAE(config, random_init).to(device)
model.apply(init_xavier)
# train
if comet:
trainer = Trainer(model, dataset, data, calc_rmse, config.epochs,
config.lr, config.weight_decay, experiment)
else:
trainer = Trainer(model, dataset, data, calc_rmse, config.epochs,
config.lr, config.weight_decay)
trainer.iterate()
elif config.optim == "SGD":
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
start_iter = 0
if config.resume:
print("Loading the trained params and the state of optimizer...")
start_iter = load_ckpt(config.resume,
[("model", model)],
[("optimizer", optimizer)])
for param_group in optimizer.param_groups:
param_group["lr"] = lr
print("Starting from iter ", start_iter)
trainer = Trainer(start_iter, config, device, model, dataset_train,
dataset_val, criterion, optimizer, experiment=experiment)
if config.comet:
with experiment.train():
trainer.iterate()
else:
trainer.iterate()
# Set the configuration for testing
elif config.mode == "test":
pass
# <model load the trained weights>
# evaluate(model, dataset_val)