Spaces:
Sleeping
Sleeping
| import torch | |
| from torch import nn as nn | |
| from torch.nn import init as init | |
| from torch.nn.modules.batchnorm import _BatchNorm | |
| def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs): | |
| if not isinstance(module_list, list): | |
| module_list = [module_list] | |
| for module in module_list: | |
| for m in module.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| init.kaiming_normal_(m.weight, **kwargs) | |
| m.weight.data *= scale | |
| if m.bias is not None: | |
| m.bias.data.fill_(bias_fill) | |
| elif isinstance(m, nn.Linear): | |
| init.kaiming_normal_(m.weight, **kwargs) | |
| m.weight.data *= scale | |
| if m.bias is not None: | |
| m.bias.data.fill_(bias_fill) | |
| elif isinstance(m, _BatchNorm): | |
| init.constant_(m.weight, 1) | |
| if m.bias is not None: | |
| m.bias.data.fill_(bias_fill) | |
| def make_layer(basic_block, num_basic_block, **kwarg): | |
| layers = [] | |
| for _ in range(num_basic_block): | |
| layers.append(basic_block(**kwarg)) | |
| return nn.Sequential(*layers) | |
| def pixel_unshuffle(x, scale): | |
| b, c, hh, hw = x.size() | |
| out_channel = c * (scale**2) | |
| h, w = hh // scale, hw // scale | |
| x_view = x.view(b, c, h, scale, w, scale) | |
| return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w) |