"""Minimal DDPM for low-dimensional tabular data with guided counterfactuals.
PEAL's image generators (DiffusionAutoencoder, PathLDM, Stable Diffusion) do
not apply to the toy tabular tasks (2-D "circle"-style datasets from
``peal.data.tabular_datasets``). This module provides an MLP score network
trained with the standard epsilon-prediction loss, plus an ``edit`` method in
the ``EditCapableGenerator`` protocol that produces counterfactuals by
classifier-guided reverse diffusion from a partially noised input. Training
writes ``diffusion.pt``, ``config.yaml`` and TensorBoard scatter plots into
``config.base_path``.
"""
import os
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm, trange
from typing import Tuple
import logging
from torch.utils.data import DataLoader
import math
from peal.generators.interfaces import EditCapableGenerator, GeneratorConfig
from peal.data.interfaces import DataConfig
import typing
from peal.log import get_logger
_log = get_logger(__name__)
[docs]
class TabularDDPMConfig(GeneratorConfig):
"""Config of :class:`TabularDDPM`.
Parameters
----------
generator_type : str
Registry name, ``"TabularDDPM"``.
input_dim : int
Number of non-label feature columns the score network denoises.
embed_dim : int
Width of the timestep embedding and of the hidden MLP layers.
num_timesteps : int
Length of the forward noising chain (``T``).
model_name : str
Filename of the saved state dict inside ``base_path``.
var_schedule : str
``"linear"`` or ``"cosine"`` beta schedule.
base_path : str
Directory used as ``model_dir`` when none is passed to the generator.
num_epochs, learning_rate : int, float
Adam training budget. The batch size comes from
``GeneratorConfig.batch_size``.
grad_scales : list of float
Classifier-gradient weights tried by :meth:`TabularDDPM.edit`.
noise_steps_for_counterfactuals : list of int
Forward-diffusion depths (``t``) tried by :meth:`TabularDDPM.edit`.
attack_iterations : int
How many times the (steps x scales) grid is repeated in ``edit``.
data : DataConfig
Dataset config used when no dataset object is passed.
"""
generator_type: str = "TabularDDPM"
input_dim: int = 2
embed_dim: int = 64
num_timesteps: int = 1000
model_name: str = "diffusion.pt"
var_schedule: str = "linear"
base_path: str = ""
num_epochs: int = 100
learning_rate: float = 1e-3
grad_scales: typing.List[float] = [1.0]
noise_steps_for_counterfactuals: typing.List[int] = [100]
attack_iterations: int = 1
data: DataConfig = DataConfig()
logging.getLogger().setLevel(logging.INFO)
[docs]
class PositionalEncoding(nn.Module):
"""Fixed sinusoidal timestep embeddings looked up by integer timestep.
Parameters
----------
embed_dim : int
Embedding width (even sizes fill sin/cos pairs exactly).
max_len : int
Number of timesteps; the table has ``max_len`` rows for steps
``0 .. max_len - 1`` (row 0 of the internal table is dropped so that
timestep 0 does not map to the all-zero/one row).
"""
[docs]
def __init__(self, embed_dim, max_len=500):
"""Precompute the ``[max_len, embed_dim]`` table as buffer ``P``."""
super(PositionalEncoding, self).__init__()
max_len += 1
P = torch.zeros(max_len, embed_dim)
freqs = torch.arange(max_len)[:, None] / (
torch.pow(
10000, torch.arange(0, embed_dim, 2, dtype=torch.float32) / embed_dim
)
)
P[:, 0::2] = torch.sin(freqs)
P[:, 1::2] = torch.cos(freqs)
self.register_buffer("P", P[1:])
[docs]
def forward(self, t):
"""Return the embeddings of timesteps ``t`` (int or LongTensor)."""
return self.P[t]
[docs]
class ScoreNetwork(nn.Module):
"""Four-layer MLP predicting the noise ``eps`` from ``x_t`` and a time embedding.
Parameters
----------
input_dim : int
Feature dimension of the data (output width).
embed_dim : int
Hidden width; the time embedding is added after the first layer.
Notes
-----
The layers are ``nn.LazyLinear`` so the input width is fixed on the first
forward pass.
"""
[docs]
def __init__(self, input_dim, embed_dim):
"""Build the lazy linear layers and the LayerNorm before the output."""
super(ScoreNetwork, self).__init__()
self.embed_dim = embed_dim
self.layer1 = nn.LazyLinear(embed_dim)
self.layer2 = nn.LazyLinear(embed_dim)
self.layer3 = nn.LazyLinear(embed_dim)
self.norm = nn.LayerNorm(embed_dim)
self.layer4 = nn.LazyLinear(input_dim)
[docs]
def forward(self, x, time_embed):
"""Predict noise for ``x`` ``[B, input_dim]`` given ``time_embed``
``[B, embed_dim]``."""
x = self.layer1(x) + time_embed
x = F.silu(self.layer2(x))
x = F.silu(self.layer3(x))
return self.layer4((self.norm(x)))
[docs]
class BasicDiscreteTimeModel(nn.Module):
"""Score network plus timestep embedding: the trainable part of the DDPM.
Parameters
----------
input_dim : int
Feature dimension of the data.
embed_dim : int
Width of the timestep embedding and MLP.
num_timesteps : int
Size of the positional-encoding table.
"""
[docs]
def __init__(self, input_dim: int, embed_dim: int, num_timesteps: int):
"""Instantiate the :class:`PositionalEncoding` and :class:`ScoreNetwork`."""
super(BasicDiscreteTimeModel, self).__init__()
self.positional_embeddings = PositionalEncoding(
embed_dim=embed_dim, max_len=num_timesteps
)
self.score_network = ScoreNetwork(input_dim=input_dim, embed_dim=embed_dim)
# self.decoder = ScoreNetwork(input_dim=input_dim, embed_dim=embed_dim)
[docs]
def forward(self, x, t):
"""Return the predicted noise for ``x`` at timestep(s) ``t``."""
time_embed = self.positional_embeddings(t)
return self.score_network(x, time_embed)
[docs]
class TabularDDPM(EditCapableGenerator):
"""DDPM generator for tabular data with classifier-guided counterfactual edits.
On construction the beta/alpha/alpha_bar schedule is registered as
buffers and :meth:`train_model` is called with ``mode="train"``, i.e.
the score network is (re)trained every time the generator is built.
Parameters
----------
config : TabularDDPMConfig
Generator config (already a pydantic object, not a path).
dataset : Dataset, optional
Tabular dataset with an ``attributes`` list; loaded from
``config.data`` via ``get_datasets`` when ``None``. Columns named
``"Confounder"`` and ``"Target"`` are excluded from the diffused
features (``input_idx``); ``"Target"`` columns form ``target_idx``.
model_dir : str, optional
Where the model and logs are written; defaults to ``config.base_path``.
device : str
Unused; tensors follow the classifier's device in ``edit``.
Attributes
----------
beta, alpha, alpha_bar : torch.Tensor
``[num_timesteps]`` schedule buffers.
model : BasicDiscreteTimeModel
The trained score network (set by :meth:`train_model`).
counterfactuals, original_sample : torch.Tensor
Last ``edit`` result and input, used by :meth:`plot_counterfactuals`.
"""
[docs]
def __init__(self, config, dataset=None, model_dir=None, device="cpu", **kwargs):
"""Resolve the dataset and schedule, then train the score network."""
super(TabularDDPM, self).__init__()
# self.config = load_yaml_config(config)
self.config = config
if not model_dir is None:
self.model_dir = model_dir
else:
# self.model_dir = config['base_path']
self.model_dir = config.base_path
# if not os.path.exists(model_dir):
# os.mkdir(model_dir)
# self.model_dir = model_dir
# self.input_dim = config['input_dim']
self.input_dim = config.input_dim
try:
# self.num_timesteps = config['num_timesteps']
self.num_timesteps = config.num_timesteps
except KeyError:
pass
if dataset is None:
from peal.data.dataset_factory import get_datasets
dataset = get_datasets(config.data)[0]
self.dataset = dataset
self.input_idx = [
idx
for idx, element in enumerate(self.dataset.attributes)
if element not in ["Confounder", "Target"]
]
self.target_idx = [
idx
for idx, element in enumerate(self.dataset.attributes)
if element == "Target"
]
# data = torch.zeros([len(dataset.data),len(dataset.attributes)], dtype=torch.float16)
# for idx, key in enumerate(dataset.data):
# data[idx] = dataset.data[key]
def schedules(num_timesteps: int, type: str = "linear"):
scale = 1000 / num_timesteps
if type == "linear":
min_var = scale * 1e-4
max_var = scale * 1e-2
return torch.linspace(
min_var, max_var, num_timesteps, dtype=torch.float32
)
elif type == "cosine":
steps = num_timesteps + 1
x = torch.linspace(0, num_timesteps, steps, dtype=torch.float64)
alphas_cumprod = (
torch.cos(
((x / num_timesteps) + scale) / (1 + scale) * math.pi * 0.5
)
** 2
)
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
return torch.clip(betas, 0, 0.999)
betas = schedules(num_timesteps=config.num_timesteps, type=config.var_schedule)
self.register_buffer("beta", betas)
self.register_buffer("alpha", 1 - self.beta)
self.register_buffer("alpha_bar", self.alpha.cumprod(0))
self.train_model(model_name=config.model_name)
[docs]
def forward_diffusion(
self, clean_x: torch.Tensor, noise: torch.tensor, timestep: torch.Tensor
):
"""Sample ``x_t = sqrt(alpha_bar_t) x_0 + sqrt(1 - alpha_bar_t) eps``.
Parameters
----------
clean_x : torch.Tensor
``[B, D]`` clean features.
noise : torch.Tensor
``[B, D]`` standard normal noise ``eps``.
timestep : int or torch.Tensor
A single timestep applied to the whole batch, or a ``[B]`` tensor
of per-sample timesteps.
Returns
-------
torch.Tensor
``[B, D]`` noised features.
"""
if isinstance(timestep, int):
timestep = torch.tensor([timestep])
alpha_bar_t = (
self.alpha_bar[timestep]
.repeat(clean_x.shape[0])[:, None]
.to(clean_x.device)
)
else:
alpha_bar_t = self.alpha_bar[timestep][:, None].to(clean_x.device)
mu = torch.sqrt(alpha_bar_t)
std = torch.sqrt(1 - alpha_bar_t)
noisy_x = mu * clean_x + std * noise
return noisy_x
[docs]
def reverse_diffusion_ddpm(
self, noisy_x: torch.Tensor, model: nn.Module, timestep: torch.Tensor
):
"""One ancestral DDPM step ``x_t -> x_{t-1}`` with variance ``beta_t``.
Parameters
----------
noisy_x : torch.Tensor
``[B, D]`` current sample ``x_t``.
model : nn.Module
Score network called as ``model(x=noisy_x, t=timestep)``.
timestep : int
Scalar timestep shared by the batch; no noise is added at ``0``.
Returns
-------
torch.Tensor
``[B, D]`` sample ``x_{t-1}``.
"""
alpha_t = (
self.alpha[timestep].repeat(noisy_x.shape[0])[:, None].to(noisy_x.device)
)
alpha_bar_t = (
self.alpha_bar[timestep]
.repeat(noisy_x.shape[0])[:, None]
.to(noisy_x.device)
)
beta_t = 1 - alpha_t
eps_hat = model(x=noisy_x, t=timestep)
posterior_mean = (1 / torch.sqrt(alpha_t)) * (
noisy_x - (beta_t / torch.sqrt(1 - alpha_bar_t) * eps_hat)
)
z = torch.randn_like(noisy_x)
if timestep > 0:
denoised_x = (
posterior_mean + torch.sqrt(beta_t) * z
) # * z * (timestep > 0)) # variance = beta_t
else:
denoised_x = posterior_mean
return denoised_x
[docs]
def train_model(self, model_name="diffusion.pt", mode="train"):
"""Train the score network or load it from ``model_dir``.
Parameters
----------
model_name : str
State-dict filename inside ``model_dir``.
mode : str
``"train"`` always trains for ``config.num_epochs`` epochs with
Adam on the epsilon-MSE loss and saves the state dict plus
``config.yaml``; any other value loads ``model_name`` if it
exists and otherwise only logs a hint.
Notes
-----
Training logs ``Loss/train`` and, every 10 epochs, a scatter plot of
generated versus real points to ``<model_dir>/logs`` (TensorBoard).
The trained/loaded network is stored as ``self.model``. The nested
``run_epoch`` helper is defined but never called.
"""
self.model_path = os.path.join(self.model_dir, model_name)
model = BasicDiscreteTimeModel(
input_dim=self.config.input_dim,
embed_dim=self.config.embed_dim,
num_timesteps=self.config.num_timesteps,
)
if not os.path.exists(self.model_dir):
os.makedirs(self.model_dir)
if model_name in os.listdir(self.model_dir) and not mode == "train":
model.load_state_dict(torch.load(self.model_path))
logging.info(f"Model found with path {self.model_path}")
elif model_name not in os.listdir(self.model_dir) and mode != "train":
logging.info(
'Model not found. Please run train_and_load_diffusion method and set its argument mode="train" '
)
else:
logging.info(f"Training model with path {self.model_path}")
def diffusion_loss(model: nn.Module, clean_x: torch.Tensor) -> torch.Tensor:
t = torch.randint(self.num_timesteps, (clean_x.shape[0],))
eps_t = torch.randn_like(clean_x)
alpha_bar_t = self.alpha_bar[t][:, None]
x_t = self.forward_diffusion(clean_x=clean_x, noise=eps_t, timestep=t)
# x_t = torch.sqrt(alpha_bar_t) * clean_x + torch.sqrt(1 - alpha_bar_t) * eps_t
eps_hat = model(x=x_t, t=t)
loss_diff = nn.MSELoss(reduction="sum")(eps_hat, eps_t)
return loss_diff
def run_epoch(
model: nn.Module, dataloader: torch.utils.data.dataloader.DataLoader
):
model.train()
epoch_loss = 0.0
for x, _ in dataloader:
optimizer.zero_grad()
loss = diffusion_loss(x)
epoch_loss += loss
loss.backward()
optimizer.step()
return epoch_loss / len(dataloader.dataset)
if mode == "train":
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
writer = SummaryWriter(os.path.join(self.model_dir, "logs"))
model.train()
# num_epochs = self.config['num_epochs']
num_epochs = self.config.num_epochs
dataloader = DataLoader(
self.dataset, batch_size=self.config.batch_size, shuffle=True
)
learning_rate = self.config.learning_rate
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
losses = []
for i in trange(num_epochs):
epoch_loss = 0.0
for x, _ in dataloader:
optimizer.zero_grad()
loss = diffusion_loss(model, x[:, self.input_idx])
epoch_loss += loss
loss.backward()
optimizer.step()
train_loss = epoch_loss / len(dataloader.dataset)
_log.info("%s", f"Epoch: {i}, train_loss: {train_loss}")
writer.add_scalar("Loss/train", train_loss, i)
losses.append(train_loss.detach().cpu().numpy())
if i % 10 == 0 or i == num_epochs - 1:
model.eval()
with torch.no_grad():
sample_x = (
self.sample_ddpm(
model, n_samples=min(1000, len(self.dataset))
)[-1]
.cpu()
.numpy()
)
real_x = (
next(iter(dataloader))[0][:, self.input_idx].cpu().numpy()
)
fig, ax = plt.subplots()
ax.scatter(
real_x[:, 0],
real_x[:, 1],
alpha=0.5,
label="Real Data",
color="blue",
)
ax.scatter(
sample_x[:, 0],
sample_x[:, 1],
alpha=0.5,
label="Generated Data",
color="orange",
)
ax.legend()
writer.add_figure("Generated Data vs Real", fig, i)
plt.close(fig)
model.train()
torch.save(model.state_dict(), self.model_path)
from peal.global_utils import save_yaml_config
save_yaml_config(self.config, os.path.join(self.model_dir, "config.yaml"))
self.model = model
@torch.no_grad()
def sample_ddpm(self, model: nn.Module, n_samples: int = 256, label=None):
"""Iteratively denoise pure noise, keeping the sample after every step.
Parameters
----------
model : nn.Module
Score network (switched to eval mode).
n_samples : int
Number of samples drawn from ``N(0, I)`` in ``input_dim`` dims.
label : optional
Unused.
Returns
-------
list of torch.Tensor
``num_timesteps + 1`` tensors of shape ``[n_samples, input_dim]``,
from the initial noise to the final sample (``[-1]``).
"""
model.eval()
x_pred = []
x = torch.randn(n_samples, self.input_dim).to(next(model.parameters()).device)
x_pred.append(x)
with torch.no_grad():
for t in reversed(range(0, self.num_timesteps)):
x = self.reverse_diffusion_ddpm(noisy_x=x, model=model, timestep=t)
x_pred.append(x)
return x_pred
[docs]
def sample_x(self, batch_size=1):
"""Return ``batch_size`` fully denoised samples ``[batch_size, input_dim]``."""
x = self.sample_ddpm(model=self.model, n_samples=batch_size)[-1]
return x
[docs]
def sample_counterfactual_ddpm(
self,
clean_batch: torch.Tensor,
model: nn.Module,
classifier: nn.Module,
num_noise_steps: int,
target_classes: int,
classifier_grad_weight: float,
):
"""Classifier-guided reverse diffusion from a partially noised batch.
The clean batch is noised to ``num_noise_steps`` and denoised again;
at every step the posterior mean is shifted by the (scaled) gradient
of the cross-entropy towards ``target_classes``, evaluated on the
fully denoised estimate of the previous step. After each guided step
the intermediate sample is denoised to ``t = 0`` without guidance to
obtain a candidate counterfactual.
Parameters
----------
clean_batch : torch.Tensor
``[B, D]`` inputs to edit.
model : nn.Module
Score network.
classifier : nn.Module
Differentiable predictor on the ``D`` features.
num_noise_steps : int
Forward-diffusion depth ``t`` to start from.
target_classes : torch.Tensor
``[B]`` target labels.
classifier_grad_weight : float
Scale of the classifier gradient.
Returns
-------
tuple
``(counterfactuals, guided_grads, unguided_grads, total_series)``:
``counterfactuals`` ``[B, num_noise_steps, D]`` (the clean input
first, then one candidate per guided step), the guided and
unconditional update terms with the same shape, and a list of the
per-step denoising trajectories. The first three are also stored
as attributes ``counterfactuals_series``, ``guided_grads`` and
``unguided_grads``.
"""
classifier.eval()
self.classifier = classifier
# DEFINE BATCH SIZE AND COUNTERFACTUAL CLASS
bs = clean_batch.shape[0]
# COMPUTE CLEAN GRADIENTS FOR THE FIRST STEP
classifier_criterion = lambda x: F.cross_entropy(classifier(x), target_classes)
clean_batch_copy = torch.nn.Parameter(clean_batch)
loss = classifier_criterion(clean_batch_copy)
loss.backward()
clean_grad = classifier_grad_weight * clean_batch_copy.grad.detach()
# REDEFINING VARIABLES AND PERFORMING FORWARD DIFFUSION
eps_t = torch.randn_like(clean_batch)
next_z = self.forward_diffusion(
clean_x=clean_batch, noise=eps_t, timestep=num_noise_steps
)
counterfactuals = [] # total counterfactuals
counterfactuals.append(clean_batch)
guided_grads = [] # guided grads at the first step
unconditional_grads = [] # diffusion grads at the first step
total_series = (
[]
) # contains evolution from noisy to cleaned instance for each data point
for i in tqdm(range(0, num_noise_steps)[::-1]):
# Denoise z_t to create z_t-1 (next z)
alpha_i = self.alpha[i].repeat(bs)[:, None].to(next_z.device)
alpha_bar_i = self.alpha_bar[i].repeat(bs)[:, None].to(next_z.device)
sigma_i = torch.sqrt(1 - self.alpha[i]).to(next_z.device)
eps_hat = model(next_z, i)
# Unconditional mean
unconditional_grad = -eps_hat / torch.sqrt(1 - alpha_bar_i)
z_t_mean = (next_z + unconditional_grad * (1 - alpha_i)) / torch.sqrt(
alpha_i
)
# Guided mean
z_t_mean -= sigma_i * (clean_grad / torch.sqrt(alpha_bar_i))
if i > 0:
next_z = z_t_mean + (sigma_i * torch.randn_like(clean_batch))
else:
next_z = z_t_mean
next_x = next_z.clone()
# Denoise to create a cleaned x (next x)
series = []
series.append(next_x.detach())
for t in range(0, i)[::-1]:
if i == 0:
break
next_x = self.reverse_diffusion_ddpm(
noisy_x=next_x, model=model, timestep=t
)
series.append(next_x.detach())
total_series.append(series)
guided_grads.append(
-sigma_i * clean_grad.detach() / torch.sqrt(alpha_bar_i)
)
unconditional_grads.append(
unconditional_grad.detach() * (1 - alpha_i) / torch.sqrt(alpha_i)
)
if i != 0:
counterfactuals.append(next_x.detach())
# Gradient wrt denoised image (next_x)
next_x_copy = torch.nn.Parameter(next_x.clone())
loss = classifier_criterion(next_x_copy)
loss.backward()
clean_classifier_grad = next_x_copy.grad.detach()
clean_grad = classifier_grad_weight * clean_classifier_grad
# self.counterfactuals = counterfactuals
# self.guided_grads = guided_grads
# self.diffusion_grads = unconditional_grads
# self.pointwise_evolution = total_series
counterfactuals = torch.stack(counterfactuals).permute(1, 0, 2)
guided_grads = torch.stack(guided_grads).permute(1, 0, 2)
unguided_grads = torch.stack(unconditional_grads).permute(1, 0, 2)
self.counterfactuals_series = counterfactuals
self.guided_grads = guided_grads
self.unguided_grads = unguided_grads
return counterfactuals, guided_grads, unguided_grads, total_series
[docs]
def discard_counterfactuals(
self,
counterfactuals,
classifier,
target_classes,
target_confidence,
minimal_counterfactuals,
tolerance=0.1,
):
"""Keep, per sample, the earliest candidate that reaches ``target_confidence``.
Parameters
----------
counterfactuals : torch.Tensor
``[B, S, D]`` candidates per sample from
:meth:`sample_counterfactual_ddpm`.
classifier : nn.Module
Predictor whose softmax confidence is checked.
target_classes : torch.Tensor
``[B]`` target labels.
target_confidence : float
Required target-class confidence.
minimal_counterfactuals : torch.Tensor
``[B, D]`` current best counterfactuals, updated in place.
tolerance : float
Unused.
Returns
-------
torch.Tensor
``minimal_counterfactuals``: unchanged rows that already satisfy
the confidence, otherwise the first satisfying candidate, or the
last candidate when none satisfies it.
"""
# compute distance of current minimal_counterefactuals from radius 1.0
# current_counterfactual_distance_from_manifold = torch.abs((torch.pow(minimal_counterfactuals, 2).sum(dim=-1) - 1.0))
for i in range(len(counterfactuals)):
# compute classifier for all the counterfactuals for each point
new_counterfactuals_confidence = classifier(counterfactuals[i]).softmax(
dim=-1
)[:, target_classes[i]]
# check if new counterfactuals satisfy the confidence constraint
new_confidence_satisfied_indices = torch.nonzero(
new_counterfactuals_confidence > target_confidence
)
current_confidence_satisfied = (
classifier(minimal_counterfactuals[i : i + 1])
.softmax(dim=-1)[0][target_classes[i]]
.item()
> target_confidence
)
# if current counterfactual satisfies confidence and tolerance, maintain status quo
if current_confidence_satisfied:
continue
elif new_confidence_satisfied_indices.nelement() != 0:
minimal_counterfactuals[i] = counterfactuals[i][
new_confidence_satisfied_indices[0].item()
]
else:
minimal_counterfactuals[i] = counterfactuals[i][-1]
return minimal_counterfactuals
[docs]
def edit(
self,
x_in: torch.Tensor,
target_confidence_goal: float,
source_classes: torch.Tensor,
target_classes: nn.Module,
classifier=None,
**kwargs,
) -> Tuple[
list[torch.Tensor],
list[torch.Tensor],
list[torch.Tensor],
list[torch.Tensor],
list[torch.Tensor],
list[torch.Tensor],
]:
"""Produce counterfactuals for ``x_in`` (``EditCapableGenerator`` protocol).
Loops ``config.attack_iterations`` times over every
``noise_steps_for_counterfactuals`` x ``grad_scales`` combination,
running :meth:`sample_counterfactual_ddpm` from the current best
counterfactuals and keeping the first candidate that reaches
``target_confidence_goal`` via :meth:`discard_counterfactuals`.
Parameters
----------
x_in : torch.Tensor
``[B, D]`` inputs (moved to the classifier's device).
target_confidence_goal : float
Required confidence for the target class.
source_classes : torch.Tensor
``[B]`` source labels (unused).
target_classes : torch.Tensor
``[B]`` target labels.
classifier : nn.Module, optional
Predictor; taken from ``kwargs["predictor"]`` when ``None``.
**kwargs
Other explainer arguments (``predictor``, ``explainer_config``,
...); only ``predictor`` is read.
Returns
-------
tuple
``(counterfactual_list, diff_latent, y_target_end_confidence,
x_list, history_list, cluster_list)``: per-sample counterfactuals,
``x_in - counterfactuals`` as a ``[B, D]`` tensor, per-sample
target-class confidences, per-sample inputs, and two lists of
empty/``None`` placeholders. The counterfactuals and inputs are
also stored as ``self.counterfactuals`` / ``self.original_sample``.
"""
if classifier is None:
classifier = kwargs.get("predictor", None)
device = next(classifier.parameters()).device
x_in = x_in.to(device)
target_classes = target_classes.to(device)
self.original_sample = x_in
# scales = self.config['grad_scales']
# noise_steps = self.config['noise_steps_for_counterfactuals']
scales = self.config.grad_scales
noise_steps = self.config.noise_steps_for_counterfactuals
num_iterations = self.config.attack_iterations
minimal_counterfactuals = x_in.clone()
for it in range(num_iterations):
for steps in noise_steps:
for s in scales:
(
counterfactuals,
guided_grads,
unguided_grads,
total_series,
) = self.sample_counterfactual_ddpm(
clean_batch=minimal_counterfactuals,
model=self.model,
classifier=classifier,
num_noise_steps=steps,
target_classes=target_classes,
classifier_grad_weight=s,
)
minimal_counterfactuals = self.discard_counterfactuals(
counterfactuals=counterfactuals,
classifier=classifier,
target_confidence=target_confidence_goal,
target_classes=target_classes,
minimal_counterfactuals=minimal_counterfactuals,
)
list_counterfactuals = [row_tensor for row_tensor in minimal_counterfactuals]
diff_latent = x_in - minimal_counterfactuals
confidences = classifier(minimal_counterfactuals).softmax(dim=-1)
y_target_end_confidence = [
confidences[i][target_classes[i]].detach()
for i in range(len(minimal_counterfactuals))
]
x_list = [row_tensor for row_tensor in x_in]
self.counterfactuals = minimal_counterfactuals
self.original_sample = x_in
history_list = [[] for _ in range(len(minimal_counterfactuals))]
cluster_list = [None for _ in range(len(minimal_counterfactuals))]
return (
list_counterfactuals,
diff_latent,
y_target_end_confidence,
x_list,
history_list,
cluster_list,
)
[docs]
def plot_counterfactuals(self):
"""Show a 2-D scatter of the dataset with arrows from each ``original_sample``
to its counterfactual (uses the first two ``input_idx`` columns)."""
plt.figure(figsize=(5, 5))
data = torch.zeros(
[len(self.dataset.data), len(self.dataset.attributes)], dtype=torch.float16
)
for idx, key in enumerate(self.dataset.data):
data[idx] = self.dataset.data[key]
_log.info("%s", data)
plt.scatter(
data[:, self.input_idx[0]],
data[:, self.input_idx[1]],
c=np.where(data[:, self.target_idx] == 0, "lightcyan", "lightgray")[0],
)
for i, point in enumerate(self.counterfactuals):
plt.scatter(
self.original_sample[i, 0], self.original_sample[i, 1], color="green"
)
plt.scatter(point[0], point[1])
plt.arrow(
self.original_sample[i, 0],
self.original_sample[i, 1],
# plot the original point plus arrow until (j+granularity)th point
point[0] - self.original_sample[i, 0],
point[1] - self.original_sample[i, 1],
head_width=0.05,
head_length=0.05,
fc="blue",
ec="blue",
)
plt.show()
__name__ = "CircleDiffusionAdaptor"