bayeshmaml.py

from copy import deepcopy

import numpy as np
import torch
from torch import nn as nn
from torch.nn import functional as F
import pytorch_lightning as pl

from methods.meta_template import MetaTemplate
import backbone
from methods.hypernets.utils import (
    get_param_dict,
    kl_diag_gauss_with_standard_gauss,
    reparameterize,
)  # __jm__ configure black and isort to use the same schema


class BHyperNet(pl.LightningModule):
    """bayesian hypernetwork for target network params"""

    def __init__(
        self, hn_hidden_size, n_way, embedding_size, feat_dim, out_neurons, params
    ):
        super().__init__()

        self.hn_head_len = params.hn_head_len

        head = [nn.Linear(embedding_size, hn_hidden_size), nn.ReLU()]

        if self.hn_head_len > 2:
            for i in range(self.hn_head_len - 2):
                head.append(nn.Linear(hn_hidden_size, hn_hidden_size))
                head.append(nn.ReLU())

        self.head = nn.Sequential(*head)

        # tails to equate weights with distributions
        tail_mean = [nn.Linear(hn_hidden_size, out_neurons)]
        tail_logvar = [nn.Linear(hn_hidden_size, out_neurons)]

        self.tail_mean = nn.Sequential(*tail_mean)
        self.tail_logvar = nn.Sequential(*tail_logvar)

    def forward(self, x):
        out = self.head(x)
        out_mean = self.tail_mean(out)
        out_logvar = self.tail_logvar(out)
        return out_mean, out_logvar


class BayesHMAML(MetaTemplate):
    def __init__(self, model_func, n_way, n_support, n_query, params, approx=False):
        super().__init__(
            model_func, n_way, n_support, n_query  # , approx=approx, params=params
        )
        self.n_query = n_query
        # loss function component
        self.loss_kld = kl_diag_gauss_with_standard_gauss  # Kullback–Leibler divergence
        self.kl_scale = params.kl_scale
        self.kl_step = None  # increase step for share of kld in loss
        self.kl_stop_val = params.kl_stop_val

        # __jm__ not defined here - either don't contribute to state or were missed out
        self.classifier = None  # __jm__ typehint later
        self.target_net_param_shapes: dict[str, torch.Size]
        self.hypernet_heads = nn.ModuleDict()

        self.delta_list = []
        self.alpha = 0

        # num of weight set draws for softvoting
        self.weight_set_num_train = params.hm_weight_set_num_train  # train phase
        self.weight_set_num_test = (
            params.hm_weight_set_num_test
            if params.hm_weight_set_num_test != 0
            else None
        )  # test phase

    def _init_classifier(self):
        assert (
            self.hn_tn_hidden_size % self.n_way == 0
        ), f"hn_tn_hidden_size {self.hn_tn_hidden_size} should be the multiple of n_way {self.n_way}"
        layers = []

        for i in range(self.hn_tn_depth):
            in_dim = self.feat_dim if i == 0 else self.hn_tn_hidden_size
            out_dim = (
                self.n_way if i == (self.hn_tn_depth - 1) else self.hn_tn_hidden_size
            )

            linear = backbone.BLinear_fw(in_dim, out_dim)
            linear.bias.data.fill_(0)

            layers.append(linear)

        self.classifier = nn.Sequential(*layers)

    def _init_hypernet_modules(self, params):
        target_net_param_dict = get_param_dict(self.classifier)

        target_net_param_dict = {
            name.replace(".", "-"): p
            # replace dots with hyphens bc torch doesn't like dots in modules names
            for name, p in target_net_param_dict.items()
        }

        self.target_net_param_shapes = {
            name: p.shape for (name, p) in target_net_param_dict.items()
        }

        self.hypernet_heads = nn.ModuleDict()

        for name, param in target_net_param_dict.items():
            if self.hm_use_class_batch_input and name[-4:] == "bias":
                # notice head_out val when using this strategy
                continue

            bias_size = param.shape[0] // self.n_way

            head_in = self.embedding_size
            head_out = (
                (param.numel() // self.n_way) + bias_size
                if self.hm_use_class_batch_input
                else param.numel()
            )
            # make hypernetwork for target network param
            self.hypernet_heads[name] = BHyperNet(
                self.hn_hidden_size,
                self.n_way,
                head_in,
                self.feat_dim,
                head_out,
                params,
            )

    def get_hn_delta_params(self, support_embeddings):
        if self.hm_detach_before_hyper_net:
            support_embeddings = support_embeddings.detach()

        if self.hm_use_class_batch_input:
            delta_params_list = []

            for name, param_net in self.hypernet_heads.items():
                support_embeddings_resh = support_embeddings.reshape(self.n_way, -1)

                delta_params_mean, params_logvar = param_net(
                    support_embeddings_resh
                )  # __jm__ this is basically forwarding through hyper
                bias_neurons_num = self.target_net_param_shapes[name][0] // self.n_way

                if self.hn_adaptation_strategy == "increasing_alpha" and self.alpha < 1:
                    delta_params_mean = delta_params_mean * self.alpha
                    params_logvar = params_logvar * self.alpha

                weights_delta_mean = (
                    delta_params_mean[:, :-bias_neurons_num]
                    .contiguous()
                    .view(*self.target_net_param_shapes[name])
                )
                bias_delta_mean = delta_params_mean[:, -bias_neurons_num:].flatten()

                weights_logvar = (
                    params_logvar[:, :-bias_neurons_num]
                    .contiguous()
                    .view(*self.target_net_param_shapes[name])
                )
                bias_logvar = params_logvar[:, -bias_neurons_num:].flatten()

                delta_params_list.append([weights_delta_mean, weights_logvar])
                delta_params_list.append([bias_delta_mean, bias_logvar])
            return delta_params_list
        else:
            delta_params_list = []

            for name, param_net in self.hypernet_heads.items():
                flattened_embeddings = support_embeddings.flatten()

                delta_mean, logvar = param_net(flattened_embeddings)

                if name in self.target_net_param_shapes.keys():
                    delta_mean = delta_mean.reshape(self.target_net_param_shapes[name])
                    logvar = logvar.reshape(self.target_net_param_shapes[name])

                if self.hn_adaptation_strategy == "increasing_alpha" and self.alpha < 1:
                    delta_mean = self.alpha * delta_mean
                    logvar = self.alpha * logvar

                delta_params_list.append([delta_mean, logvar])
            return delta_params_list

    def _update_weight(self, weight, update_mean, logvar, train_stage=False):
        """get distribution associated with weight. Sample weights for target network."""
        if update_mean is None and logvar is None:
            return
        # if weight.mu is None:
        if not hasattr(weight, "mu") or weight.mu is None:
            weight.mu = None  # __jm__ what the fuck
            weight.mu = weight - update_mean
        else:
            weight.mu = weight.mu - update_mean

        if logvar is None:  # used in maml warmup
            weight.fast = []
            weight.fast.append(weight.mu)
        else:
            weight.logvar = logvar

            weight.fast = []
            if train_stage:
                for _ in range(
                    self.weight_set_num_train
                ):  # sample fast parameters for training
                    weight.fast.append(reparameterize(weight.mu, weight.logvar))
            else:
                if self.weight_set_num_test is not None:
                    for _ in range(
                        self.weight_set_num_test
                    ):  # sample fast parameters for testing
                        weight.fast.append(reparameterize(weight.mu, weight.logvar))
                else:
                    weight.fast.append(weight.mu)  # return expected value

    def _scale_step(self):
        """calculate regularization step for kld"""
        if self.kl_step is None:
            # scale step is calculated so that share of kld in loss increases kl_scale -> kl_stop_val
            self.kl_step = np.power(
                1 / self.kl_scale * self.kl_stop_val, 1 / self.stop_epoch
            )

        self.kl_scale = self.kl_scale * self.kl_step

    def _get_p_value(self):
        if self.epoch < self.hm_maml_warmup_epochs:
            return 1.0

        elif (
            self.hm_maml_warmup_epochs
            <= self.epoch
            < self.hm_maml_warmup_epochs + self.hm_maml_warmup_switch_epochs
        ):
            return (
                self.hm_maml_warmup_switch_epochs
                + self.hm_maml_warmup_epochs
                - self.epoch
            ) / (self.hm_maml_warmup_switch_epochs + 1)
        return 0.0

    def _update_network_weights(
        self,
        delta_params_list,
        support_embeddings,
        support_data_labels,
        train_stage=False,
    ):
        if self.hm_maml_warmup and not self.single_test:
            p = self._get_p_value()
            # warmup coef p decreases 1 -> 0
            if p > 0.0:
                fast_parameters = []

                clf_fast_parameters = list(self.classifier.parameters())
                for weight in self.classifier.parameters():
                    weight.fast = None
                    weight.mu = None
                    # weight.logvar = None
                self.classifier.zero_grad()
                fast_parameters = (
                    fast_parameters + clf_fast_parameters
                )  # __jm__ whyyyy are you adding them together

                for _task_step in range(self.task_update_num):
                    scores = self.classifier(support_embeddings)

                    set_loss = self.loss_fn(scores, support_data_labels)
                    reduction = self.kl_scale
                    for weight in self.classifier.parameters():
                        if weight.logvar is not None:
                            if weight.mu is not None:
                                # set_loss = set_loss + self.kl_w * reduction * self.loss_kld(weight.mu, weight.logvar)
                                set_loss = set_loss + reduction * self.loss_kld(
                                    weight.mu, weight.logvar
                                )
                            else:
                                # set_loss = set_loss + self.kl_w * reduction * self.loss_kld(weight, weight.logvar)
                                set_loss = set_loss + reduction * self.loss_kld(
                                    weight, weight.logvar
                                )

                    grad = torch.autograd.grad(
                        set_loss, fast_parameters, create_graph=True, allow_unused=True
                    )  # build full graph support gradient of gradient

                    if self.approx:
                        grad = [
                            g.detach() for g in grad
                        ]  # do not calculate gradient of gradient if using first order approximation

                    if p == 1:
                        # update weights of classifier network by adding gradient
                        for k, weight in enumerate(self.classifier.parameters()):
                            update_value = self.train_lr * grad[k]
                            update_mean, logvar = delta_params_list[k]
                            self._update_weight(
                                weight, update_value, logvar, train_stage
                            )

                    elif 0.0 < p < 1.0:
                        # update weights of classifier network by adding gradient and output of hypernetwork
                        for k, weight in enumerate(self.classifier.parameters()):
                            update_value = self.train_lr * p * grad[k]
                            update_mean, logvar = delta_params_list[k]
                            update_mean = (1 - p) * update_mean + update_value
                            self._update_weight(
                                weight, update_mean, logvar, train_stage
                            )
            else:
                for k, weight in enumerate(self.classifier.parameters()):
                    update_mean, logvar = delta_params_list[k]
                    self._update_weight(weight, update_mean, logvar, train_stage)
        else:
            for k, weight in enumerate(self.classifier.parameters()):
                update_mean, logvar = delta_params_list[k]
                self._update_weight(weight, update_mean, logvar, train_stage)

    def _get_list_of_delta_params(
        self, _maml_warmup_used, support_embeddings, support_data_labels
    ):
        # if not maml_warmup_used:

        if self.enhance_embeddings:
            with torch.no_grad():
                logits = self.classifier.forward(support_embeddings).detach()
                logits = F.softmax(logits, dim=1)

            labels = support_data_labels.view(support_embeddings.shape[0], -1)
            support_embeddings = torch.cat((support_embeddings, logits, labels), dim=1)

        for weight in self.parameters():
            weight.fast = None
        for weight in self.classifier.parameters():
            weight.mu = None
            # weight.logvar = None
        self.zero_grad()

        support_embeddings = self.apply_embeddings_strategy(support_embeddings)

        delta_params = self.get_hn_delta_params(support_embeddings)

        if self.hm_save_delta_params and len(self.delta_list) == 0:
            self.delta_list = [{"delta_params": delta_params}]

        return delta_params

    def set_forward_loss(self, x):
        """Adapt and forward using x. Return scores and total losses"""
        scores, total_delta_sum = self.set_forward(
            x, is_feature=False  # , train_stage=True
        )

        # calc_sigma = calc_sigma and (self.epoch == self.stop_epoch - 1 or self.epoch % 100 == 0)
        # sigma, mu = self._mu_sigma(calc_sigma)

        query_data_labels = torch.repeat_interleave(range(self.n_way), self.n_query)
        if self.hm_support_set_loss:
            support_data_labels = torch.repeat_interleave(
                range(self.n_way), self.n_support
            )
            query_data_labels = torch.cat((support_data_labels, query_data_labels))

        reduction = self.kl_scale

        loss_ce = self.loss_fn(scores, query_data_labels)

        loss_kld = torch.zeros_like(loss_ce)

        for _name, weight in self.classifier.named_parameters():
            if weight.mu is not None and weight.logvar is not None:
                val = self.loss_kld(weight.mu, weight.logvar)
                # loss_kld = loss_kld + self.kl_w * reduction * val
                loss_kld = loss_kld + reduction * val

        loss = loss_ce + loss_kld

        if self.hm_lambda != 0:
            loss = loss + self.hm_lambda * total_delta_sum

        _topk_scores, topk_labels = scores.data.topk(1, 1, True, True)
        topk_ind = topk_labels.cpu().numpy().flatten()
        y_labels = query_data_labels.cpu().numpy()
        top1_correct = np.sum(topk_ind == y_labels)
        task_accuracy = (top1_correct / len(query_data_labels)) * 100

        return loss, loss_ce, loss_kld, task_accuracy

    def set_forward_loss_with_adaptation(self, x):
        """returns loss and accuracy from adapted model (copy)"""
        scores, _ = self.set_forward(
            x, is_feature=False  # , train_stage=False
        )  # scores from adapted copy
        support_data_labels = torch.repeat_interleave(range(self.n_way), self.n_support)

        reduction = self.kl_scale

        loss_ce = self.loss_fn(scores, support_data_labels)

        loss_kld = torch.zeros_like(loss_ce)

        for _name, weight in self.classifier.named_parameters():
            if weight.mu is not None and weight.logvar is not None:
                # loss_kld = loss_kld + self.kl_w * reduction * self.loss_kld(weight.mu, weight.logvar)
                loss_kld = loss_kld + reduction * self.loss_kld(
                    weight.mu, weight.logvar
                )

        loss = loss_ce + loss_kld

        _topk_scores, topk_labels = scores.data.topk(1, 1, True, True)
        topk_ind = topk_labels.cpu().numpy().flatten()
        y_labels = support_data_labels.cpu().numpy()
        top1_correct = np.sum(topk_ind == y_labels)
        task_accuracy = (top1_correct / len(support_data_labels)) * 100

        return loss, task_accuracy

    def train_loop(self, _epoch, train_loader, optimizer):  # overwrite parrent function
        print_freq = 10
        avg_loss = 0
        task_count = 0
        loss_all = []
        loss_ce_all = []
        loss_kld_all = []
        # loss_kld_no_scale_all = []
        acc_all = []
        optimizer.zero_grad()

        # train
        for i, (x, _) in enumerate(train_loader):
            self.n_query = x.size(1) - self.n_support
            assert self.n_way == x.size(0), "MAML does not support way change"

            loss, loss_ce, loss_kld, task_accuracy = self.set_forward_loss(x)
            avg_loss = avg_loss + loss.item()  # .data[0]
            loss_all.append(loss)
            loss_ce_all.append(loss_ce.item())
            loss_kld_all.append(loss_kld.item())
            # loss_kld_no_scale_all.append(loss_kld_no_scale.item())
            acc_all.append(task_accuracy)

            task_count += 1

            if task_count == self.n_task:  # MAML update several tasks at one time
                loss_q = torch.stack(loss_all).sum(0)
                loss_q.backward()

                optimizer.step()
                task_count = 0
                loss_all = []

            optimizer.zero_grad()
            if i % print_freq == 0:
                print(
                    "Epoch {:d}/{:d} | Batch {:d}/{:d} | Loss {:f}".format(
                        self.epoch,
                        self.stop_epoch,
                        i,
                        len(train_loader),
                        avg_loss / float(i + 1),
                    )
                )

        self._scale_step()

        acc_all = np.asarray(acc_all)
        acc_mean = np.mean(acc_all)

        metrics = {"accuracy/train": acc_mean}

        loss_ce_all = np.asarray(loss_ce_all)
        loss_ce_mean = np.mean(loss_ce_all)

        metrics["loss_ce"] = loss_ce_mean

        loss_kld_all = np.asarray(loss_kld_all)
        loss_kld_mean = np.mean(loss_kld_all)

        metrics["loss_kld"] = loss_kld_mean

        if self.hn_adaptation_strategy == "increasing_alpha":
            metrics["alpha"] = self.alpha

        if self.hm_save_delta_params and len(self.delta_list) > 0:
            delta_params = {"epoch": self.epoch, "delta_list": self.delta_list}
            metrics["delta_params"] = delta_params

        if self.alpha < 1:
            self.alpha += self.hn_alpha_step

        return metrics

    def set_forward_with_adaptation(self, x: torch.Tensor):
        self_copy = deepcopy(self)

        # deepcopy does not copy "fast" parameters so it should be done manually
        for param1, param2 in zip(
            self.feature.parameters(), self_copy.feature.parameters()
        ):
            if hasattr(param1, "fast"):
                if param1.fast is not None:
                    param2.fast = param1.fast.clone()
                else:
                    param2.fast = None

        for param1, param2 in zip(
            self.classifier.parameters(), self_copy.classifier.parameters()
        ):
            if hasattr(param1, "fast"):
                if param1.fast is not None:
                    param2.fast = list(param1.fast)
                else:
                    param2.fast = None
            if hasattr(param1, "mu"):
                if param1.mu is not None:
                    param2.mu = param1.mu.clone()
                else:
                    param2.mu = None
            if hasattr(param1, "logvar"):
                if param1.logvar is not None:
                    param2.logvar = param1.logvar.clone()
                else:
                    param2.logvar = None

        metrics = {"accuracy/val@-0": self_copy.query_accuracy(x)}

        val_opt_type = (
            torch.optim.Adam if self.hn_val_optim == "adam" else torch.optim.SGD
        )
        val_opt = val_opt_type(self_copy.parameters(), lr=self.hn_val_lr)

        if self.hn_val_epochs > 0:
            for i in range(1, self.hn_val_epochs + 1):
                self_copy.train()
                val_opt.zero_grad()
                loss, val_support_acc = self_copy.set_forward_loss_with_adaptation(x)
                loss.backward()
                val_opt.step()
                self_copy.eval()
                metrics[f"accuracy/val_support_acc@-{i}"] = val_support_acc
                metrics[f"accuracy/val_loss@-{i}"] = loss.item()
                metrics[f"accuracy/val@-{i}"] = self_copy.query_accuracy(x)

        # free CUDA memory by deleting "fast" parameters
        for param in self_copy.parameters():
            param.fast = None
            param.mu = None
            param.logvar = None

        return metrics[f"accuracy/val@-{self.hn_val_epochs}"], metrics