| import torch |
|
|
| def get_sample_align_fn(sample_align_model): |
| r""" |
| Code is adapted from https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/scripts/classifier_sample.py#L54-L61 |
| """ |
| def sample_align_fn(x, *args, **kwargs): |
| r""" |
| Calculates `grad(log(p(y|x)))` |
| This uses the conditioning strategy from Sohl-Dickstein et al. (2015). |
| |
| Parameters |
| ---------- |
| x: torch.Tensor |
| |
| Returns |
| ------- |
| grad |
| """ |
| |
| with torch.enable_grad(): |
| x_in = x.detach().requires_grad_(True) |
| logits = sample_align_model(x_in, *args, **kwargs) |
| grad = torch.autograd.grad(logits.sum(), x_in, allow_unused=True)[0] |
| return grad |
| return sample_align_fn |
|
|
| def get_alignment_kwargs_avg_x(context_seq=None, target_seq=None, ): |
| r""" |
| Please customize this function for generating knowledge "avg_x_gt" |
| that guides the inference. |
| E.g., this function uses 2.0 ground-truth future average intensity as "avg_x_gt" for demonstration. |
| |
| Parameters |
| ---------- |
| context_seq: torch.Tensor, aka "y" |
| target_seq: torch.Tensor, aka "x" |
| |
| Returns |
| ------- |
| alignment_kwargs: Dict |
| """ |
| multiplier = 2.0 |
| batch_size = target_seq.shape[0] |
| ret = torch.mean(target_seq.view(batch_size, -1), |
| dim=1, keepdim=True) * multiplier |
| return {"avg_x_gt": ret} |
