diff options
| -rw-r--r-- | experiments/resnet_crossover_smoke.py | 80 | ||||
| -rw-r--r-- | sdil/conv_crossover.py | 147 |
2 files changed, 227 insertions, 0 deletions
diff --git a/experiments/resnet_crossover_smoke.py b/experiments/resnet_crossover_smoke.py index 8cdcb0c..6d18c7f 100644 --- a/experiments/resnet_crossover_smoke.py +++ b/experiments/resnet_crossover_smoke.py @@ -10,6 +10,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from sdil.conv import CIFARLocalResNet from sdil.conv_crossover import ( CIFARDualPropResNet, + CIFAREquilibriumPropResNet, CIFARForwardForwardResNet, CIFARPEPITAResNet, ) @@ -249,9 +250,88 @@ def audit_forward_forward(): } +def audit_equilibrium_propagation(): + common = dict( + depth=8, base_width=2, seed=101, normalization="batchnorm", + residual_scale=1.0, dtype=torch.float64) + reference = CIFARLocalResNet(**common) + net = CIFAREquilibriumPropResNet( + **common, ep_beta=0.5, dt=0.5, free_steps=8, + nudge_steps=2, random_beta_sign=False) + generator = torch.Generator().manual_seed(102) + image = torch.randn(3, 3, 32, 32, generator=generator, + dtype=torch.float64) + labels = torch.tensor([0, 5, 9]) + one_hot = torch.nn.functional.one_hot(labels, 10).to(torch.float64) + with torch.no_grad(): + free = net.ep_settle( + image, one_hot, beta=0.0, steps=net.ep_free_steps) + signed_beta = net.ep_beta + nudged = net.ep_settle( + image, one_hot, beta=signed_beta, steps=net.ep_nudge_steps, + initial_states=free) + parameters = ( + net.W + net.gamma + net.beta + [net.W_out, net.b_out]) + for parameter in parameters: + parameter.requires_grad_(True) + + def phase_phi(states): + clipped_image = net.ep_rho(image) + hidden = [net.ep_rho(value.detach()) for value in states[:-1]] + value = image.new_zeros(()) + for index in range(net.n_hidden): + prediction, _ = net._node_prediction( + index, hidden, clipped_image) + value += torch.sum( + prediction * net.ep_rho(states[index].detach())) + features = hidden[-1].mean(dim=(2, 3)) + output_prediction = features @ net.W_out.t() + net.b_out + value += torch.sum( + output_prediction * net.ep_rho(states[-1].detach())) + return value + + objective = ( + phase_phi(free) - phase_phi(nudged) + ) / (signed_beta * image.shape[0]) + gradients = torch.autograd.grad(objective, parameters) + (directions, gamma_directions, beta_directions, + output_weight, output_bias) = net.ep_ascent_directions( + image, free, nudged, signed_beta) + actual = ( + directions + gamma_directions + beta_directions + + [output_weight, output_bias]) + errors = [ + relative_error(direction, -gradient) + for direction, gradient in zip(actual, gradients)] + for parameter in parameters: + parameter.requires_grad_(False) + state_bounds = [ + (float(value.min()), float(value.max())) + for value in free + nudged] + if (net.n_forward_parameters != reference.n_forward_parameters + or max(errors) >= 2e-12 + or any(low < 0 or high > 1 for low, high in state_bounds) + or not all(torch.isfinite(value).all() + for value in free + nudged)): + raise AssertionError({ + "direction_errors": errors, + "state_bounds": state_bounds, + }) + return { + "contrastive_direction_max_relative_error": max(errors), + "matched_forward_parameter_count": net.n_forward_parameters, + "free_steps": net.ep_free_steps, + "nudge_steps": net.ep_nudge_steps, + "state_minimum": min(low for low, _ in state_bounds), + "state_maximum": max(high for _, high in state_bounds), + "uses_reverse_task_loss_graph": False, + } + + def main(): print(json.dumps({ "dualprop": audit_dualprop(), + "equilibrium_propagation": audit_equilibrium_propagation(), "forward_forward": audit_forward_forward(), "pepita": audit_pepita(), }, diff --git a/sdil/conv_crossover.py b/sdil/conv_crossover.py index 092bffe..7b34ab7 100644 --- a/sdil/conv_crossover.py +++ b/sdil/conv_crossover.py @@ -478,3 +478,150 @@ class CIFARDualPropResNet(CIFARLocalResNet): gamma_directions=gamma_directions, beta_directions=beta_directions) return float(loss) + + +class CIFAREquilibriumPropResNet(CIFARDualPropResNet): + """Canonical two-phase EP dynamics on the matched residual energy graph. + + The Dual-Propagation parent class is used only for its audited residual-edge + operators and exact edge transposes. EP itself has one hard-sigmoid state + per forward population, synchronous leaky relaxation, random nudge sign, + and free-versus-nudged local correlations. + """ + + def __init__(self, *args, ep_beta=0.5, dt=0.5, free_steps=20, + nudge_steps=4, random_beta_sign=True, **kwargs): + super().__init__( + *args, alpha=0.0, dp_beta=0.1, inference_passes=1, **kwargs) + if ep_beta <= 0 or not 0 < dt <= 1: + raise ValueError("invalid Equilibrium Propagation scale or step") + if free_steps < 1 or nudge_steps < 1: + raise ValueError("EP phases require positive relaxation steps") + self.ep_beta = float(ep_beta) + self.ep_dt = float(dt) + self.ep_free_steps = int(free_steps) + self.ep_nudge_steps = int(nudge_steps) + self.ep_random_beta_sign = bool(random_beta_sign) + + @staticmethod + def ep_rho(state): + return state.clamp(0.0, 1.0) + + @staticmethod + def ep_rhop(state): + return ((state >= 0.0) & (state <= 1.0)).to(state.dtype) + + @torch.no_grad() + def ep_settle(self, image, one_hot, beta=0.0, steps=None, + initial_states=None): + steps = self.ep_free_steps if steps is None else int(steps) + clipped_image = self.ep_rho(image) + if initial_states is None: + shapes = self.forward( + clipped_image, training=True, + update_stats=False)["hiddens"] + states = [torch.zeros_like(value) for value in shapes] + states.append(torch.zeros( + image.shape[0], self.n_classes, + device=image.device, dtype=image.dtype)) + else: + states = [value.detach().clone() for value in initial_states] + for _ in range(steps): + rho_hidden = [ + self.ep_rho(value) for value in states[:-1]] + fields = rho_hidden + [self.ep_rho(states[-1])] + updated = [] + for index, state in enumerate(states[:-1]): + prediction, _ = self._node_prediction( + index, rho_hidden, clipped_image) + feedback = self._outgoing_feedback( + index, fields, rho_hidden, clipped_image) + drive = -self.ep_rho(state) + prediction + feedback + updated.append(self.ep_rho( + state + self.ep_dt * self.ep_rhop(state) * drive)) + features = rho_hidden[-1].mean(dim=(2, 3)) + prediction = features @ self.W_out.t() + self.b_out + output = states[-1] + drive = -self.ep_rho(output) + prediction + if beta: + drive = drive + 2.0 * beta * (one_hot - output) + updated.append(self.ep_rho( + output + self.ep_dt * self.ep_rhop(output) * drive)) + states = updated + return states + + @torch.no_grad() + def _ep_phase_correlations(self, image, states): + clipped_image = self.ep_rho(image) + rho_hidden = [self.ep_rho(value) for value in states[:-1]] + batch = image.shape[0] + directions = [] + gamma_directions = [] + beta_directions = [] + for index in range(self.n_hidden): + _, cache = self._node_prediction( + index, rho_hidden, clipped_image) + spec = self.layer_specs[index] + field = self.ep_rho(states[index]) * spec.branch_scale + field, gamma_direction, beta_direction = ( + self._normalization_backward( + index, field, cache["normalization"])) + direction = torch.nn.grad.conv2d_weight( + cache["pre"], self.W[index].shape, field, + stride=spec.stride, padding=spec.padding) + directions.append(direction / batch) + if gamma_direction is not None: + gamma_directions.append(gamma_direction / batch) + beta_directions.append(beta_direction / batch) + features = rho_hidden[-1].mean(dim=(2, 3)) + output = self.ep_rho(states[-1]) + output_weight = output.t() @ features / batch + output_bias = output.mean(dim=0) + return ( + directions, gamma_directions, beta_directions, + output_weight, output_bias) + + @torch.no_grad() + def ep_ascent_directions(self, image, free_states, nudged_states, + signed_beta): + if signed_beta == 0: + raise ValueError("EP contrast requires a nonzero nudge") + free = self._ep_phase_correlations(image, free_states) + nudged = self._ep_phase_correlations(image, nudged_states) + result = [] + for nudged_group, free_group in zip(nudged, free): + if isinstance(nudged_group, list): + result.append([ + (nudged_value - free_value) / signed_beta + for nudged_value, free_value in zip( + nudged_group, free_group)]) + else: + result.append((nudged_group - free_group) / signed_beta) + return tuple(result) + + def ep_step(self, image, labels, eta, eta_output=None, + momentum=0.0, weight_decay=0.0, generator=None): + one_hot = F.one_hot(labels, self.n_classes).to(image.dtype) + with torch.no_grad(): + free_states = self.ep_settle( + image, one_hot, beta=0.0, steps=self.ep_free_steps) + sign = 1.0 + if self.ep_random_beta_sign: + draw = torch.randint( + 0, 2, (), device=image.device, generator=generator) + sign = 1.0 if int(draw) else -1.0 + signed_beta = sign * self.ep_beta + nudged_states = self.ep_settle( + image, one_hot, beta=signed_beta, + steps=self.ep_nudge_steps, initial_states=free_states) + (directions, gamma_directions, beta_directions, + output_weight, output_bias) = self.ep_ascent_directions( + image, free_states, nudged_states, signed_beta) + self.apply_ascent( + directions, output_weight, output_bias, eta, + eta_output=eta_output, momentum=momentum, + weight_decay=weight_decay, + gamma_directions=gamma_directions, + beta_directions=beta_directions) + loss = F.mse_loss(free_states[-1], one_hot) + return float(loss), free_states |