| import torch, sys, json |
| sys.path.insert(0, "..") |
| from src.moo import cagrad_update_k2, raco_update_k2, _grads_dot, _solve_lambda_k2 |
|
|
| def test_clipping_triggers(): |
| """When CAGrad wants a large correction on the less-preferred objective, |
| RACO should clip it.""" |
| w1, w2 = 0.8, 0.2 |
| torch.manual_seed(42) |
|
|
| g1 = torch.randn(100) * 3.0 |
| g2 = -g1 * 0.9 + torch.randn(100) * 0.3 |
|
|
| g1_list = [g1] |
| g2_list = [g2] |
|
|
| cagrad_g = cagrad_update_k2(g1_list, g2_list, w1, w2, c=0.5) |
| raco_g = raco_update_k2(g1_list, g2_list, w1, w2, c=0.5) |
|
|
| cg = cagrad_g[0] |
| rg = raco_g[0] |
|
|
| alignment_ca_g1 = float((cg @ g1).item()) |
| alignment_ca_g2 = float((cg @ g2).item()) |
| alignment_ra_g1 = float((rg @ g1).item()) |
| alignment_ra_g2 = float((rg @ g2).item()) |
|
|
| result = { |
| "test": "Strong conflicting objectives, w1=0.8, w2=0.2", |
| "w1": w1, "w2": w2, |
| "g1_norm": float(g1.norm().item()), |
| "g2_norm": float(g2.norm().item()), |
| "cos_angle": float((g1 @ g2 / (g1.norm() * g2.norm())).item()), |
| "cagrad_alignment_g1": alignment_ca_g1, |
| "cagrad_alignment_g2": alignment_ca_g2, |
| "raco_alignment_g1": alignment_ra_g1, |
| "raco_alignment_g2": alignment_ra_g2, |
| "cagrad_more_g2_magnitude": abs(alignment_ca_g2) > abs(alignment_ra_g2), |
| "raco_better_respects_weights": alignment_ra_g2 > alignment_ca_g2, |
| } |
| return result |
|
|
| def test_theorem_convergence(): |
| """Test that CAGrad-Clip provides valid descent directions.""" |
| results = {} |
| for c in [0.1, 0.3, 0.5, 0.7]: |
| w1 = 0.7 |
| w2 = 0.3 |
| torch.manual_seed(123) |
| g1 = torch.randn(50) * 2.0 |
| g2 = -g1 * 0.7 + torch.randn(50) * 0.5 |
| g1_list, g2_list = [g1], [g2] |
|
|
| try: |
| rg = raco_update_k2(g1_list, g2_list, w1, w2, c=c) |
| g_update = rg[0] |
| g0 = g1 * w1 + g2 * w2 |
|
|
| descent_ok = float((g_update @ g_update).item()) > 0 |
| results[f"c={c}"] = { |
| "valid": descent_ok, |
| "g0_norm": float(g0.norm().item()), |
| "g_update_norm": float(g_update.norm().item()), |
| } |
| except Exception as e: |
| results[f"c={c}"] = {"error": str(e)} |
| return results |
|
|
| if __name__ == "__main__": |
| all_tests = {} |
| all_tests["test_clipping"] = test_clipping_triggers() |
| all_tests["test_convergence"] = test_theorem_convergence() |
|
|
| print(json.dumps(all_tests, indent=2)) |
|
|
| with open("/Users/equan_p/Developer/playground/ICML-2/repro_raco/outputs/tier1/test_results.json", "w") as f: |
| json.dump(all_tests, f, indent=2) |
|
|
| all_pass = all( |
| t.get("raco_better_respects_weights", True) or t.get("valid", True) |
| for t in all_tests.values() |
| ) |
| print(f"\nAll tests passed: {all_pass}") |
|
|
