def __init__(self, checkpoint, config):
self.config = config
# Initialize the network
self.model = models.MuZeroNetwork(self.config)
self.model.set_weights(checkpoint["weights"])
self.model.eval()
def __init__(self, initial_checkpoint, config):
self.config = config
# Fix random generator seed
numpy.random.seed(self.config.seed)
torch.manual_seed(self.config.seed)
# Initialize the network
self.model = models.MuZeroNetwork(self.config)
self.model.set_weights(initial_checkpoint["weights"])
self.model.to(torch.device("cuda" if self.config.reanalyse_on_gpu else "cpu"))
self.model.eval()
self.num_reanalysed_games = initial_checkpoint["num_reanalysed_games"]
def __init__(self, initial_checkpoint, Game, config, seed):
self.config = config
self.game = Game(seed)
# Fix random generator seed
numpy.random.seed(seed)
torch.manual_seed(seed)
# Initialize the network
self.model = models.MuZeroNetwork(self.config)
self.model.set_weights(initial_checkpoint["weights"])
self.model.to(
torch.device("cuda" if torch.cuda.is_available() else "cpu"))
self.model.eval()
def __init__(self, initial_checkpoint, config):
self.config = config
# Fix random generator seed
numpy.random.seed(self.config.seed)
torch.manual_seed(self.config.seed)
# Initialize the network
self.model = models.MuZeroNetwork(self.config)
self.model.set_weights(copy.deepcopy(initial_checkpoint["weights"]))
self.model.to(torch.device("cuda" if self.config.train_on_gpu else "cpu"))
self.model.train()
self.training_step = initial_checkpoint["training_step"]
if "cuda" not in str(next(self.model.parameters()).device):
print("You are not training on GPU.\n")
# Initialize the optimizer
if self.config.optimizer == "SGD":
self.optimizer = torch.optim.SGD(
self.model.parameters(),
lr=self.config.lr_init,
momentum=self.config.momentum,
weight_decay=self.config.weight_decay,
)
elif self.config.optimizer == "Adam":
self.optimizer = torch.optim.Adam(
self.model.parameters(),
lr=self.config.lr_init,
weight_decay=self.config.weight_decay,
)
else:
raise NotImplementedError(
f"{self.config.optimizer} is not implemented. You can change the optimizer manually in trainer.py."
)
if initial_checkpoint["optimizer_state"] is not None:
print("Loading optimizer...\n")
self.optimizer.load_state_dict(
copy.deepcopy(initial_checkpoint["optimizer_state"])
)
def get_initial_weights(self, config):
model = models.MuZeroNetwork(config)
weigths = model.get_weights()
summary = str(model).replace("\n", " \n\n")
return weigths, summary