| |
| |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def rescale_2x(x: torch.Tensor, scale_factor=2): |
| return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False) |
|
|
|
|
| def resize_to(x: torch.Tensor, tgt_hw: tuple): |
| return F.interpolate(x, size=tgt_hw, mode="bilinear", align_corners=False) |
|
|
|
|
| def clip_grad(params, mode, clip_cfg: dict): |
| if mode == "norm": |
| if "max_norm" not in clip_cfg: |
| raise ValueError("`clip_cfg` must contain `max_norm`.") |
| torch.nn.utils.clip_grad_norm_( |
| params, |
| max_norm=clip_cfg.get("max_norm"), |
| norm_type=clip_cfg.get("norm_type", 2.0), |
| ) |
| elif mode == "value": |
| if "clip_value" not in clip_cfg: |
| raise ValueError("`clip_cfg` must contain `clip_value`.") |
| torch.nn.utils.clip_grad_value_(params, clip_value=clip_cfg.get("clip_value")) |
| else: |
| raise NotImplementedError |
|
|
