diff options
Diffstat (limited to 'sdil/conv_crossover.py')
| -rw-r--r-- | sdil/conv_crossover.py | 147 |
1 files changed, 147 insertions, 0 deletions
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 |