File size: 3,642 Bytes
9f29df6 | 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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """PyTorch implementation of RainNet v1.0."""
import torch
from torch import nn
class RainNet(nn.Module):
"""The 20-convolution encoder-decoder from Ayzel et al. (2020)."""
def __init__(self, in_channels=4, out_channels=1, dropout=0.5, mode="regression"):
super().__init__()
if mode != "regression":
raise ValueError("This package implements RainNet regression mode only")
self.in_channels = in_channels
self.pool = nn.MaxPool2d(2)
self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
self.relu = nn.ReLU(inplace=True)
self.conv1f = nn.Conv2d(in_channels, 64, 3, padding=1)
self.conv1s = nn.Conv2d(64, 64, 3, padding=1)
self.conv2f = nn.Conv2d(64, 128, 3, padding=1)
self.conv2s = nn.Conv2d(128, 128, 3, padding=1)
self.conv3f = nn.Conv2d(128, 256, 3, padding=1)
self.conv3s = nn.Conv2d(256, 256, 3, padding=1)
self.conv4f = nn.Conv2d(256, 512, 3, padding=1)
self.conv4s = nn.Conv2d(512, 512, 3, padding=1)
self.drop4 = nn.Dropout(dropout)
self.conv5f = nn.Conv2d(512, 1024, 3, padding=1)
self.conv5s = nn.Conv2d(1024, 1024, 3, padding=1)
self.drop5 = nn.Dropout(dropout)
self.conv6f = nn.Conv2d(1536, 512, 3, padding=1)
self.conv6s = nn.Conv2d(512, 512, 3, padding=1)
self.conv7f = nn.Conv2d(768, 256, 3, padding=1)
self.conv7s = nn.Conv2d(256, 256, 3, padding=1)
self.conv8f = nn.Conv2d(384, 128, 3, padding=1)
self.conv8s = nn.Conv2d(128, 128, 3, padding=1)
self.conv9f = nn.Conv2d(192, 64, 3, padding=1)
self.conv9s = nn.Conv2d(64, 64, 3, padding=1)
self.conv9out = nn.Conv2d(64, 2, 3, padding=1)
self.output = nn.Conv2d(2, out_channels, 1)
self.apply(self._initialize)
@staticmethod
def _initialize(module):
if isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu")
if module.bias is not None:
nn.init.zeros_(module.bias)
def forward(self, x):
if x.ndim != 4:
raise ValueError(f"RainNet expects BCHW input, got shape {tuple(x.shape)}")
if x.shape[1] != self.in_channels:
raise ValueError(f"RainNet expects {self.in_channels} channels, got {x.shape[1]}")
if x.shape[-2] % 16 or x.shape[-1] % 16:
raise ValueError(
f"Spatial dimensions must be divisible by 16, got {tuple(x.shape[-2:])}"
)
c1 = self.relu(self.conv1s(self.relu(self.conv1f(x))))
c2 = self.relu(self.conv2s(self.relu(self.conv2f(self.pool(c1)))))
c3 = self.relu(self.conv3s(self.relu(self.conv3f(self.pool(c2)))))
c4 = self.relu(self.conv4s(self.relu(self.conv4f(self.pool(c3)))))
c5 = self.relu(self.conv5s(self.relu(self.conv5f(self.pool(self.drop4(c4))))))
c5 = self.drop5(c5)
c6 = torch.cat((self.upsample(c5), c4), dim=1)
c6 = self.relu(self.conv6s(self.relu(self.conv6f(c6))))
c7 = torch.cat((self.upsample(c6), c3), dim=1)
c7 = self.relu(self.conv7s(self.relu(self.conv7f(c7))))
c8 = torch.cat((self.upsample(c7), c2), dim=1)
c8 = self.relu(self.conv8s(self.relu(self.conv8f(c8))))
c9 = torch.cat((self.upsample(c8), c1), dim=1)
c9 = self.relu(self.conv9s(self.relu(self.conv9f(c9))))
return self.output(self.relu(self.conv9out(c9)))
def build_rainnet(in_channels=4, out_channels=1, dropout=0.5, mode="regression"):
return RainNet(in_channels, out_channels, dropout, mode)
|