diff options
Diffstat (limited to 'sdil/conv_crossover.py')
| -rw-r--r-- | sdil/conv_crossover.py | 175 |
1 files changed, 175 insertions, 0 deletions
diff --git a/sdil/conv_crossover.py b/sdil/conv_crossover.py index 329ad3b..092bffe 100644 --- a/sdil/conv_crossover.py +++ b/sdil/conv_crossover.py @@ -102,6 +102,181 @@ class CIFARPEPITAResNet(CIFARLocalResNet): return float(loss) +class CIFARForwardForwardResNet(CIFARLocalResNet): + """Greedy supervised Forward-Forward on the residual parameter topology.""" + + def __init__(self, *args, threshold=2.0, learning_rate=0.03, + score_from_layer=1, **kwargs): + super().__init__(*args, **kwargs) + if learning_rate <= 0 or threshold <= 0: + raise ValueError("invalid Forward-Forward hyperparameters") + self.ff_threshold = float(threshold) + self.ff_score_from_layer = int(score_from_layer) + self.ff_num_layers = self.n_hidden + 1 + if not 0 <= self.ff_score_from_layer < self.ff_num_layers: + raise ValueError("Forward-Forward score range excludes all layers") + self.ff_optimizers = [] + for index, weight in enumerate(self.W): + parameters = [weight] + if self.normalization == "batchnorm": + parameters.extend([self.gamma[index], self.beta[index]]) + for parameter in parameters: + parameter.requires_grad_(True) + self.ff_optimizers.append( + torch.optim.Adam(parameters, lr=learning_rate)) + self.W_out.requires_grad_(True) + self.b_out.requires_grad_(True) + self.ff_optimizers.append(torch.optim.Adam( + [self.W_out, self.b_out], lr=learning_rate)) + + @staticmethod + def _ff_normalize(value): + axes = tuple(range(1, value.ndim)) + norm = torch.sqrt(torch.sum(value.square(), dim=axes, keepdim=True)) + return value / (norm + 1e-8) + + def ff_overlay(self, image, labels): + flat = image.reshape(image.shape[0], -1).clone() + flat[:, :self.n_classes] = 0.0 + flat[torch.arange(image.shape[0], device=image.device), labels] = ( + torch.max(image).detach()) + return flat.reshape_as(image) + + def ff_forward(self, image, training=False, update_stats=False): + """Forward candidate-labelled data through normalized residual edges.""" + hiddens = [] + caches = [] + pre = self._ff_normalize(image) + convolution = F.conv2d(pre, self.W[0], stride=1, padding=1) + normalized, norm_cache = self._normalize( + 0, convolution, training, update_stats) + hidden = F.relu(normalized) + hiddens.append(hidden) + caches.append({ + "pre": pre, "normalization": norm_cache, "shortcut": None}) + for block in self.blocks: + first = block["first"] + second = block["second"] + parent = hidden + pre = self._ff_normalize(parent) + convolution = F.conv2d( + pre, self.W[first], stride=block["stride"], padding=1) + normalized, norm_cache = self._normalize( + first, convolution, training, update_stats) + first_hidden = F.relu(normalized) + hiddens.append(first_hidden) + caches.append({ + "pre": pre, "normalization": norm_cache, "shortcut": None}) + + pre = self._ff_normalize(first_hidden) + shortcut = self._option_a_shortcut( + parent, block["out_channels"], block["stride"]) + convolution = F.conv2d( + pre, self.W[second], stride=1, padding=1) + normalized, norm_cache = self._normalize( + second, convolution, training, update_stats) + hidden = F.relu( + shortcut + self.residual_scale * normalized) + hiddens.append(hidden) + caches.append({ + "pre": pre, "normalization": norm_cache, + "shortcut": shortcut}) + features = self._ff_normalize(hidden.mean(dim=(2, 3))) + logits = features @ self.W_out.t() + self.b_out + return { + "hiddens": hiddens, + "caches": caches, + "features": features, + "logits": logits, + } + + def _ff_local_outputs(self, layer_index, positive, negative): + """Re-evaluate exactly one target layer on detached prefix states.""" + if layer_index == self.n_hidden: + inputs = torch.cat( + [positive["features"], negative["features"]], dim=0).detach() + outputs = inputs @ self.W_out.t() + self.b_out + else: + positive_cache = positive["caches"][layer_index] + negative_cache = negative["caches"][layer_index] + inputs = torch.cat( + [positive_cache["pre"], negative_cache["pre"]], + dim=0).detach() + spec = self.layer_specs[layer_index] + convolution = F.conv2d( + inputs, self.W[layer_index], + stride=spec.stride, padding=spec.padding) + normalized, _ = self._normalize( + layer_index, convolution, training=True, update_stats=True) + if positive_cache["shortcut"] is None: + outputs = F.relu(normalized) + else: + shortcut = torch.cat([ + positive_cache["shortcut"], + negative_cache["shortcut"], + ], dim=0).detach() + outputs = F.relu( + shortcut + spec.branch_scale * normalized) + return outputs.chunk(2, dim=0) + + def ff_train_layer(self, layer_index, image, labels, + learning_rate=0.03, negative_labels=None): + """Update one greedy FF layer; every prefix and non-target is detached.""" + if not 0 <= layer_index < self.ff_num_layers: + raise ValueError("invalid Forward-Forward layer") + if negative_labels is None: + offsets = torch.randint( + 1, self.n_classes, labels.shape, device=labels.device) + negative_labels = (labels + offsets) % self.n_classes + with torch.no_grad(): + positive = self.ff_forward( + self.ff_overlay(image, labels), + training=True, update_stats=False) + negative = self.ff_forward( + self.ff_overlay(image, negative_labels), + training=True, update_stats=False) + all_parameters = ( + self.W + self.gamma + self.beta + [self.W_out, self.b_out]) + for parameter in all_parameters: + parameter.grad = None + optimizer = self.ff_optimizers[layer_index] + optimizer.param_groups[0]["lr"] = learning_rate + positive_output, negative_output = self._ff_local_outputs( + layer_index, positive, negative) + axes = tuple(range(1, positive_output.ndim)) + positive_goodness = positive_output.square().mean(dim=axes) + negative_goodness = negative_output.square().mean(dim=axes) + loss = ( + F.softplus(-positive_goodness + self.ff_threshold) + + F.softplus(negative_goodness - self.ff_threshold) + ).mean() + loss.backward() + optimizer.step() + return { + "loss": float(loss.detach()), + "positive_goodness": float(positive_goodness.mean().detach()), + "negative_goodness": float(negative_goodness.mean().detach()), + "pair_accuracy": float( + (positive_goodness > negative_goodness).float().mean()), + } + + @torch.no_grad() + def ff_candidate_scores(self, image): + scores = [] + for candidate in range(self.n_classes): + labels = torch.full( + (image.shape[0],), candidate, device=image.device, + dtype=torch.long) + forward = self.ff_forward( + self.ff_overlay(image, labels), training=False) + layer_goodness = [ + value.square().mean(dim=tuple(range(1, value.ndim))) + for value in forward["hiddens"] + [forward["logits"]]] + scores.append(sum( + layer_goodness[self.ff_score_from_layer:])) + return torch.stack(scores, dim=1) + + class CIFARDualPropResNet(CIFARLocalResNet): """Dual Propagation on the residual DAG with author DP-transpose updates. |