| | |
| | |
| |
|
| | 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 |
| |
|
