diff options
Diffstat (limited to 'experiments')
| -rw-r--r-- | experiments/two_state_debias_smoke.py | 22 |
1 files changed, 22 insertions, 0 deletions
diff --git a/experiments/two_state_debias_smoke.py b/experiments/two_state_debias_smoke.py index 9b57819..fc659e6 100644 --- a/experiments/two_state_debias_smoke.py +++ b/experiments/two_state_debias_smoke.py @@ -12,6 +12,7 @@ ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from sdil.two_state_debias import ( # noqa: E402 + BatchedLocalAffineDebiaser, LocalAffineDebiaser, two_state_difference, ) @@ -105,6 +106,24 @@ def main() -> None: else: raise AssertionError("requires-grad input was not rejected") + feature = torch.linspace(-1.0, 1.0, 64).reshape(64, 1, 1) + feature = torch.cat((feature, feature.square()), dim=2) + target = 0.2 + torch.tensor([[[0.7, -0.4]]]) * feature + batched_affine = BatchedLocalAffineDebiaser( + [feature], feature_centers=[0.0], feature_scales=[1.0], affine=True) + batched_constant = BatchedLocalAffineDebiaser( + [feature], feature_centers=[0.0], feature_scales=[1.0], affine=False) + for _ in range(20): + batched_affine.update_neutral([feature], [target], 0.2) + batched_constant.update_neutral([feature], [target], 0.2) + held_feature = torch.tensor([[[-1.3, 1.4]], [[1.3, 1.4]]]) + held_target = 0.2 + torch.tensor([[[0.7, -0.4]]]) * held_feature + held_affine = batched_affine.residual([held_feature], [held_target])[0] + held_constant = batched_constant.residual([held_feature], [held_target])[0] + batched_affine_rmse = float(held_affine.square().mean().sqrt()) + batched_constant_rmse = float(held_constant.square().mean().sqrt()) + assert batched_affine_rmse < 0.05 * batched_constant_rmse + print({ "affine_heldout_rmse": float(affine_error), "constant_heldout_rmse": float(constant_error), @@ -112,6 +131,9 @@ def main() -> None: "common_mode_max_float_error": common_mode_error, "downstream_independence_exact": True, "requires_grad_rejected": True, + "batched_affine_heldout_rmse": batched_affine_rmse, + "batched_constant_heldout_rmse": batched_constant_rmse, + "batched_neutral_observations": batched_affine.neutral_observations, }) |