1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
|
From 9fc597c20d6b758ce6b0c3a326b8bc1b9262b0d0 Mon Sep 17 00:00:00 2001
From: YurenHao0426 <Blackhao0426@gmail.com>
Date: Mon, 27 Jul 2026 13:25:10 -0500
Subject: [PATCH 09/19] crossover: isolate formal validation from test
---
config/cli_config.py | 10 ++++++++++
train.py | 43 ++++++++++++++++++++++++++-----------------
train_ff.py | 22 +++++++++++++++-------
3 files changed, 51 insertions(+), 24 deletions(-)
diff --git a/config/cli_config.py b/config/cli_config.py
index 4979b37..67a04fe 100644
--- a/config/cli_config.py
+++ b/config/cli_config.py
@@ -100,6 +100,16 @@ parser.add_argument(
'validation evaluation. The author-compatible default is full; '
'use none for diagnostic-free timed crossover runs.'))
+parser.add_argument(
+ '--test-policy', default='author', choices=['author', 'none'],
+ help=('author restores the best-validation checkpoint and evaluates test; '
+ 'none keeps test completely untouched for formal validation runs.'))
+
+parser.add_argument(
+ '--early-stop-policy', default='author', choices=['author', 'none'],
+ help=('author retains the original severe-validation-drop stop; none runs '
+ 'the full schedule unless a nonfinite training loss occurs.'))
+
dtypes = {'bfloat16': jnp.bfloat16, 'float16':jnp.float16, 'float32':jnp.float32}
parser.add_argument('--dtype', default='float32', choices=dtypes.keys())
parser.add_argument('--param-dtype', default='float32', choices=['bfloat16', 'float16', 'float32'])
diff --git a/train.py b/train.py
index b1df1b0..0b23304 100644
--- a/train.py
+++ b/train.py
@@ -43,7 +43,7 @@ for experiment_index, seed in enumerate(config.seeds):
print(msg)
loginfo_and_print(f"\n\tStarting experiment {experiment_index+1}/{len(config.seeds)}. Current seed is {seed}")
- loginfo_and_print(f"\model settings:\n{config.model}")
+ loginfo_and_print(f"\nmodel settings:\n{config.model}")
rng = jax.random.PRNGKey(seed)
rng, init_rng = jax.random.split(rng)
# Define model
@@ -179,23 +179,32 @@ for experiment_index, seed in enumerate(config.seeds):
loginfo_and_print("NaN or Inf encountered. Terminating training loop early")
break
if (epoch>5) and (hist["val_accuracy"][epoch-1] < 0.5*best_accuracy):
- loginfo_and_print("Terminating training early as validation accuracy has drastically dropped")
- break
-
- loginfo_and_print(f"\n====Loading model with best validation accuracy (epoch {best_epoch})====")
- best_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)
- rng, input_rng = jax.random.split(rng)
- if config.learning_algorithm == "ep":
- test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_ep_model(
- best_state, best_state.params, config.test_ds, config.batch_size,
- config.num_classes, config.ep_free_steps, config.ep_dt)
+ if config.early_stop_policy == "author":
+ loginfo_and_print("Terminating training early as validation accuracy has drastically dropped")
+ break
+
+ hist["best_validation_accuracy"] = best_accuracy
+ hist["best_epoch"] = best_epoch
+ hist["epochs_completed"] = epoch
+ if config.test_policy == "author":
+ loginfo_and_print(f"\n====Loading model with best validation accuracy (epoch {best_epoch})====")
+ best_state = checkpoints.restore_checkpoint(ckpt_dir=CKPT_DIR, target=state)
+ rng, input_rng = jax.random.split(rng)
+ if config.learning_algorithm == "ep":
+ test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_ep_model(
+ best_state, best_state.params, config.test_ds,
+ config.batch_size, config.num_classes, config.ep_free_steps,
+ config.ep_dt)
+ else:
+ test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(
+ best_state, best_state.params, config.test_ds,
+ config.batch_size, config.num_classes, input_rng,
+ spectral_diagnostics=(
+ config.spectral_diagnostics == "full"))
+ hist['test_loss'], hist['test_accuracy'], hist['test_top5accuracy'], hist['test_time'] = test_loss, test_accuracy, test_top5accuracy, test_time
+ loginfo_and_print('test: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (test_loss, test_accuracy, test_top5accuracy, test_time))
else:
- test_loss, test_accuracy, test_top5accuracy, test_time, _, _, _, _ = eval_model(
- best_state, best_state.params, config.test_ds, config.batch_size,
- config.num_classes, input_rng,
- spectral_diagnostics=(config.spectral_diagnostics == "full"))
- hist['test_loss'], hist['test_accuracy'], hist['test_top5accuracy'], hist['test_time'] = test_loss, test_accuracy, test_top5accuracy, test_time
- loginfo_and_print('test: \tloss: %.4f, \taccuracy: %.4f, \ttop5_accuracy: %.4f, \truntime: %.4f' % (test_loss, test_accuracy, test_top5accuracy, test_time))
+ loginfo_and_print("\n====Test evaluation disabled by validation-only policy====")
if (config.learning_algorithm != "backprop"
and config.gradient_diagnostics == "full"):
diff --git a/train_ff.py b/train_ff.py
index 3839dd3..ada45a3 100644
--- a/train_ff.py
+++ b/train_ff.py
@@ -75,9 +75,12 @@ for experiment_index, seed in enumerate(config.seeds):
validation_accuracy, validation_time = eval_ff_model(
state, config.val_ds, config.batch_size, config.num_classes,
config.ff_score_from_layer)
- test_accuracy, test_time = eval_ff_model(
- state, config.test_ds, config.batch_size, config.num_classes,
- config.ff_score_from_layer)
+ test_accuracy = np.nan
+ test_time = np.nan
+ if config.test_policy == "author":
+ test_accuracy, test_time = eval_ff_model(
+ state, config.test_ds, config.batch_size, config.num_classes,
+ config.ff_score_from_layer)
history["final"] = {
"validation_accuracy": validation_accuracy,
"validation_time": validation_time,
@@ -85,8 +88,13 @@ for experiment_index, seed in enumerate(config.seeds):
"test_time": test_time,
"train_and_eval_wall": time.time() - started,
}
- print(
- f"final val_accuracy={validation_accuracy:.3f}% "
- f"test_accuracy={test_accuracy:.3f}%",
- flush=True)
+ if config.test_policy == "author":
+ print(
+ f"final val_accuracy={validation_accuracy:.3f}% "
+ f"test_accuracy={test_accuracy:.3f}%",
+ flush=True)
+ else:
+ print(
+ f"final val_accuracy={validation_accuracy:.3f}% test=disabled",
+ flush=True)
np.save(os.path.join(outpath, "hist.npy"), history)
--
2.54.0
|
