Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn as nn | |
| from torchvision.ops.misc import FrozenBatchNorm2d | |
| def freeze_batch_norm_2d(module, module_match={}, name=""): | |
| """ | |
| Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is | |
| itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and | |
| returned. Otherwise, the module is walked recursively and submodules are converted in place. | |
| Args: | |
| module (torch.nn.Module): Any PyTorch module. | |
| module_match (dict): Dictionary of full module names to freeze (all if empty) | |
| name (str): Full module name (prefix) | |
| Returns: | |
| torch.nn.Module: Resulting module | |
| Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762 | |
| """ | |
| res = module | |
| is_match = True | |
| if module_match: | |
| is_match = name in module_match | |
| if is_match and isinstance( | |
| module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm) | |
| ): | |
| res = FrozenBatchNorm2d(module.num_features) | |
| res.num_features = module.num_features | |
| res.affine = module.affine | |
| if module.affine: | |
| res.weight.data = module.weight.data.clone().detach() | |
| res.bias.data = module.bias.data.clone().detach() | |
| res.running_mean.data = module.running_mean.data | |
| res.running_var.data = module.running_var.data | |
| res.eps = module.eps | |
| else: | |
| for child_name, child in module.named_children(): | |
| full_child_name = ".".join([name, child_name]) if name else child_name | |
| new_child = freeze_batch_norm_2d(child, module_match, full_child_name) | |
| if new_child is not child: | |
| res.add_module(child_name, new_child) | |
| return res | |
| def do_mixup(x, mixup_lambda): | |
| """ | |
| Args: | |
| x: (batch_size , ...) | |
| mixup_lambda: (batch_size,) | |
| Returns: | |
| out: (batch_size, ...) | |
| """ | |
| out = ( | |
| x.transpose(0, -1) * mixup_lambda | |
| + torch.flip(x, dims=[0]).transpose(0, -1) * (1 - mixup_lambda) | |
| ).transpose(0, -1) | |
| return out | |
| def interpolate(x, ratio): | |
| """Interpolate data in time domain. This is used to compensate the | |
| resolution reduction in downsampling of a CNN. | |
| Args: | |
| x: (batch_size, time_steps, classes_num) | |
| ratio: int, ratio to interpolate | |
| Returns: | |
| upsampled: (batch_size, time_steps * ratio, classes_num) | |
| """ | |
| (batch_size, time_steps, classes_num) = x.shape | |
| upsampled = x[:, :, None, :].repeat(1, 1, ratio, 1) | |
| upsampled = upsampled.reshape(batch_size, time_steps * ratio, classes_num) | |
| return upsampled | |
| def pad_framewise_output(framewise_output, frames_num): | |
| """Pad framewise_output to the same length as input frames. The pad value | |
| is the same as the value of the last frame. | |
| Args: | |
| framewise_output: (batch_size, frames_num, classes_num) | |
| frames_num: int, number of frames to pad | |
| Outputs: | |
| output: (batch_size, frames_num, classes_num) | |
| """ | |
| pad = framewise_output[:, -1:, :].repeat( | |
| 1, frames_num - framewise_output.shape[1], 1 | |
| ) | |
| """tensor for padding""" | |
| output = torch.cat((framewise_output, pad), dim=1) | |
| """(batch_size, frames_num, classes_num)""" | |