File size: 1,923 Bytes
2a13495 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import torch
import torch.nn as nn
def patch_first_conv(model, new_in_channels, default_in_channels=3, pretrained=True):
"""Change first convolution layer input channels.
In case:
in_channels == 1 or in_channels == 2 -> reuse original weights
in_channels > 3 -> make random kaiming normal initialization
"""
# get first conv
for module in model.modules():
if isinstance(module, nn.Conv2d) and module.in_channels == default_in_channels:
break
weight = module.weight.detach()
module.in_channels = new_in_channels
if not pretrained:
module.weight = nn.parameter.Parameter(
torch.Tensor(
module.out_channels,
new_in_channels // module.groups,
*module.kernel_size
)
)
module.reset_parameters()
elif new_in_channels == 1:
new_weight = weight.sum(1, keepdim=True)
module.weight = nn.parameter.Parameter(new_weight)
else:
new_weight = torch.Tensor(
module.out_channels, new_in_channels // module.groups, *module.kernel_size
)
for i in range(new_in_channels):
new_weight[:, i] = weight[:, i % default_in_channels]
new_weight = new_weight * (default_in_channels / new_in_channels)
module.weight = nn.parameter.Parameter(new_weight)
def replace_strides_with_dilation(module, dilation_rate):
"""Patch Conv2d modules replacing strides with dilation"""
for mod in module.modules():
if isinstance(mod, nn.Conv2d):
mod.stride = (1, 1)
mod.dilation = (dilation_rate, dilation_rate)
kh, kw = mod.kernel_size
mod.padding = ((kh // 2) * dilation_rate, (kh // 2) * dilation_rate)
# Kostyl for EfficientNet
if hasattr(mod, "static_padding"):
mod.static_padding = nn.Identity()
|