Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mistake] #1297

Closed
wants to merge 1 commit into from
Closed

[mistake] #1297

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion examples/alphazero/train.py → examples/alphazero/clean_train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@
class Config(BaseModel):
env_id: pgx.EnvId = "go_9x9"
seed: int = 0
max_num_iters: int = 400
max_num_iters: int = 500
# network params
num_channels: int = 128
num_layers: int = 6
Expand All @@ -53,6 +53,13 @@ class Config(BaseModel):
learning_rate: float = 0.001
# eval params
eval_interval: int = 5
# logging
host_name: str = os.uname().nodename
file_name: str = os.path.basename(__file__) if '__file__' in globals() else ipynbname.name()+".ipynb"
num_devices:int = len(jax.local_devices())
device_kind:str = jax.local_devices()[0].device_kind



class Config:
extra = "forbid"
Expand Down Expand Up @@ -243,6 +250,7 @@ def body_fn(val):
(my_logits, _), _ = forward.apply(
my_model_params, my_model_state, state.observation, is_eval=True
)
my_logits = 10000 * my_logits
opp_logits, _ = baseline(state.observation)
is_my_turn = (state.current_player == my_player).reshape((-1, 1))
logits = jnp.where(is_my_turn, my_logits, opp_logits)
Expand Down Expand Up @@ -296,6 +304,8 @@ def body_fn(val):
f"eval/vs_baseline/lose_rate": ((R == -1).sum() / R.size).item(),
}
)
log["score/vs_baseline_100"] = 0.5 * R.mean().item() + 0.5


# Store checkpoints
model_0, opt_state_0 = jax.tree_util.tree_map(lambda x: x[0], (model, opt_state))
Expand Down Expand Up @@ -354,6 +364,8 @@ def body_fn(val):

et = time.time()
hours += (et - st) / 3600

# Logging
log.update(
{
"train/policy_loss": policy_loss,
Expand All @@ -362,3 +374,12 @@ def body_fn(val):
"frames": frames,
}
)
sim_play = config.selfplay_batch_size * 1 * max_num_steps * iteration
sim_plan = config.selfplay_batch_size * num_simulations * max_num_steps * iteration
log["cost/simulator_evaluations/planning"] = sim_plan
log["cost/simulator_evaluations/playing"] = sim_play
log["cost/simulator_evaluations/total"] = sim_plan + sim_play
log["cost/simulator_evaluations/total [million]"] = (sim_plan + sim_play) / (10**6)
log["cost/hours/total"] = hours
log["cost/gpu_hours/total"] = hours * num_devices

Loading