SirStone
0b17430735
Save ACAdamStates (m/v tensors + t counters) as .npy files alongside network weights in latest/ and checkpoint dirs; save round counter to round_counter.txt. loadBestAvailable restores both on startup; fresh start works unchanged when files are absent. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
106 lines
3.3 KiB
Nim
106 lines
3.3 KiB
Nim
## test_weights.nim — assert-based tests for weights.nim
|
|
## Run: nim c tests/test_weights.nim && ./tests/test_weights
|
|
|
|
import std/[os, math]
|
|
import arraymancer
|
|
import "../network"
|
|
import "../training"
|
|
import "../weights"
|
|
|
|
template check(cond: bool, msg: string) =
|
|
if not cond:
|
|
quit("FAIL: " & msg, 1)
|
|
|
|
const tmpBase = "/tmp/test_weights_nim"
|
|
|
|
# ── saveWeights / loadWeights roundtrip ───────────────────────────────────────
|
|
|
|
block testRoundtrip:
|
|
let dir = tmpBase & "_roundtrip"
|
|
removeDir(dir)
|
|
|
|
let ac1 = initActorCritic()
|
|
saveWeights(ac1, dir)
|
|
|
|
var ac2 = initActorCritic()
|
|
loadWeights(ac2, dir)
|
|
|
|
# Verify a sample of tensors
|
|
template tensorEq(a, b: Tensor[float32]) =
|
|
check a.shape == b.shape, "shape mismatch"
|
|
let diff = abs(a - b)
|
|
var maxDiff = 0.0'f32
|
|
for v in diff: maxDiff = max(maxDiff, v)
|
|
check maxDiff < 1e-6'f32, "tensor values differ by " & $maxDiff
|
|
|
|
tensorEq(ac1.actor.w1, ac2.actor.w1)
|
|
tensorEq(ac1.actor.b1, ac2.actor.b1)
|
|
tensorEq(ac1.actor.w3, ac2.actor.w3)
|
|
tensorEq(ac1.critic.w1, ac2.critic.w1)
|
|
tensorEq(ac1.critic.b3, ac2.critic.b3)
|
|
tensorEq(ac1.logStd, ac2.logStd)
|
|
|
|
removeDir(dir)
|
|
|
|
# ── saveWeightsAtomic ─────────────────────────────────────────────────────────
|
|
|
|
block testAtomic:
|
|
let dir = tmpBase & "_atomic"
|
|
removeDir(dir)
|
|
|
|
let ac = initActorCritic()
|
|
saveWeightsAtomic(ac, dir)
|
|
|
|
check dirExists(dir), "targetDir should exist after atomic save"
|
|
for f in ["actor_w1.npy", "critic_w1.npy", "log_std.npy"]:
|
|
check fileExists(dir / f), "missing file: " & f
|
|
|
|
removeDir(dir)
|
|
|
|
# ── loadBestAvailable ─────────────────────────────────────────────────────────
|
|
|
|
block testLoadBest:
|
|
let root = tmpBase & "_loadbest"
|
|
removeDir(root)
|
|
createDir(root)
|
|
|
|
let ac0 = initActorCritic()
|
|
|
|
# Try with no weights — should return false
|
|
var acEmpty = initActorCritic()
|
|
var adamEmpty: ACAdamStates
|
|
check not loadBestAvailable(acEmpty, adamEmpty, root).loaded, "should return false with no weights"
|
|
|
|
# Save to latest/; should load
|
|
saveWeights(ac0, root / "latest")
|
|
var ac1 = initActorCritic()
|
|
var adam1: ACAdamStates
|
|
check loadBestAvailable(ac1, adam1, root).loaded, "should load from latest/"
|
|
|
|
# Remove latest/, save to checkpoint_1/ — should fall back
|
|
removeDir(root / "latest")
|
|
saveWeights(ac0, root / "checkpoint_1")
|
|
var ac2 = initActorCritic()
|
|
var adam2: ACAdamStates
|
|
check loadBestAvailable(ac2, adam2, root).loaded, "should load from checkpoint_1/"
|
|
|
|
removeDir(root)
|
|
|
|
# ── cleanStaleTempDirs ────────────────────────────────────────────────────────
|
|
|
|
block testClean:
|
|
let root = tmpBase & "_clean"
|
|
removeDir(root)
|
|
createDir(root)
|
|
|
|
let stale = root / "latest_tmp_12345"
|
|
createDir(stale)
|
|
check dirExists(stale), "stale dir should exist before clean"
|
|
|
|
cleanStaleTempDirs(root)
|
|
check not dirExists(stale), "stale dir should be gone after clean"
|
|
|
|
removeDir(root)
|
|
|
|
echo "All tests passed"
|