File size: 4,640 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | # Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
from typing import Dict, Iterable, Protocol
import torch
from torch_scatter import scatter
from ...diffusion.corruption.multi_corruption import MultiCorruption
from ...diffusion.data.batched_data import BatchedData
from ...diffusion.score_models.base import Diffusable
class Metric(Protocol):
"""
Computes a metric to be logged during training.
Each metric must have a name which is used as a prefix for the metric in the log.
"""
name: str
def __call__(
self,
*,
loss_per_sample_per_field: Dict[str, torch.Tensor],
multi_corruption: MultiCorruption,
score_model_output: Diffusable,
t: torch.Tensor,
batch_idx: Dict[str, torch.LongTensor],
batch: BatchedData,
noisy_batch: BatchedData,
) -> Dict[str, torch.Tensor]:
"""
Computes a metric to be logged during training. Useful, e.g., for plotting loss over time.
Args:
loss_per_sample_per_field: Dict[str, torch.Tensor], where each tensor has shape (batch_size,).
multi_corruption: MultiCorruption
score_model_output: the output produced by the model per field.
t: shape (batch_size,). Time for each element in the loss.
batch_idx: Dict[str, torch.LongTensor]: batch indices per field
batch: BatchedData: the clean (un-perturbed) batched data
noisy_batch: BatchedData: the corrupted batched data
"""
pass
def loss_per_time_bin(
loss_per_sample: torch.Tensor, t: torch.Tensor, bins: torch.Tensor
) -> torch.Tensor:
"""
Aggregate loss per bin. Useful for plotting loss over time.
Args:
loss_per_sample: shape (batch_size,). Loss for each sample.
t: shape (batch_size,). Time for each element in the loss.
bins: shape (num_bins,). Upper boundaries of the time bins.
Returns:
avg_loss_per_bin: shape (num_bins,). Average loss per time bin.
"""
bin_per_element = torch.bucketize(t, bins)
avg_loss_per_bin = scatter(
src=loss_per_sample, index=bin_per_element, dim_size=bins.shape[0], reduce="mean"
)
return avg_loss_per_bin
class LossPerTimeBin(Metric):
name = "loss_per_time_bin"
def __init__(self, t_min: float = 0.0, t_max: float = 1.0, num_bins: int = 10):
self.bins = torch.linspace(t_min, t_max, num_bins + 1)
def __call__(
self,
*,
loss_per_sample_per_field: Dict[str, torch.Tensor],
t: torch.Tensor,
**_,
) -> Dict[str, torch.Tensor]:
"""
Compute loss bins per diffusion time bin. Useful for plotting loss over diffusion time.
"""
metrics_dict = {}
for k, v in loss_per_sample_per_field.items():
assert v.shape == t.shape
# first bin is always empty because no time is less than t_min, so we skip it
avg_loss_per_bin = loss_per_time_bin(
loss_per_sample_per_field[k],
t,
bins=self.bins.to(loss_per_sample_per_field[k].device)[1:],
)
metrics_dict.update(
{
f"{k}_{self.bins[ix]:.2f}-{self.bins[ix + 1]:.2f}": avg_loss_per_bin[ix]
for ix in range(len(avg_loss_per_bin))
if avg_loss_per_bin[ix] > 0.0
}
)
return metrics_dict
class MetricsCalculator:
"""
Computes a set of metrics to be logged during training.
"""
def __init__(self, metric_fns: Iterable[Metric]):
self.metric_fns = metric_fns
def __call__(
self,
*,
loss_per_sample_per_field: Dict[str, torch.Tensor],
multi_corruption: MultiCorruption,
score_model_output: torch.Tensor,
t: torch.Tensor,
batch_idx: Dict[str, torch.LongTensor],
batch: BatchedData,
noisy_batch: BatchedData,
) -> Dict[str, torch.Tensor]:
metrics_dict = {}
for metric_fn in self.metric_fns:
_metrics_dict = metric_fn(
loss_per_sample_per_field=loss_per_sample_per_field,
multi_corruption=multi_corruption,
score_model_output=score_model_output,
t=t,
batch_idx=batch_idx,
batch=batch,
noisy_batch=noisy_batch,
)
# prepend metric name to each metric
metrics_dict.update({f"{metric_fn.name}_{k}": v for k, v in _metrics_dict.items()})
return metrics_dict
|