BADSparticles.py

from copy import deepcopy

import torch
import torch.nn as nn
import torch.nn.functional as F


class BADSParticle(nn.Module):
    def __init__(self, model_template, n_samples):
        super().__init__()
        self.model = model_template
        self.global_w = nn.Parameter(torch.randn(n_samples), requires_grad=True)

    def get_flat_param(self):
        params = []
        for p in self.model.parameters():
            params.append(p.view(-1))
        params.append(self.global_w.view(-1))
        return torch.cat(params)

    def get_flat_grad(self):
        grads = []
        for p in self.model.parameters():
            if p.grad is not None:
                grads.append(p.grad.view(-1))
            else:
                grads.append(torch.zeros_like(p).view(-1))
        if self.global_w.grad is not None:
            grads.append(self.global_w.grad.view(-1))
        else:
            grads.append(torch.zeros_like(self.global_w).view(-1))
        return torch.cat(grads)

    def set_from_vector(self, vec):
        offset = 0
        for p in self.model.parameters():
            numel = p.numel()
            p.data.copy_(vec[offset:offset+numel].view_as(p))
            offset += numel
        # set global_w
        numel = self.global_w.numel()
        self.global_w.data.copy_(vec[offset:offset+numel].view_as(self.global_w))

    def zero_grad(self):
        for p in self.model.parameters():
            if p.grad is not None:
                p.grad.zero_()
        if self.global_w.grad is not None:
            self.global_w.grad.zero_()