File size: 1,105 Bytes
d2f661a | 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 | import torch
from torch import nn
from torch.nn import functional as F
from ..nowcast.nowcast import AFNONowcastNetBase
from ..blocks.resnet import ResBlock3D
class AFNONowcastNetCascade(AFNONowcastNetBase):
def __init__(self, *args, cascade_depth=4, **kwargs):
super().__init__(*args, **kwargs)
self.cascade_depth = cascade_depth
self.resnet = nn.ModuleList()
ch = self.embed_dim_out
self.cascade_dims = [ch]
for i in range(cascade_depth-1):
ch_out = 2*ch
self.cascade_dims.append(ch_out)
self.resnet.append(
ResBlock3D(ch, ch_out, kernel_size=(1,3,3), norm=None)
)
ch = ch_out
def forward(self, x):
x = super().forward(x)
img_shape = tuple(x.shape[-2:])
cascade = {img_shape: x}
for i in range(self.cascade_depth-1):
x = F.avg_pool3d(x, (1,2,2))
x = self.resnet[i](x)
img_shape = tuple(x.shape[-2:])
cascade[img_shape] = x
return cascade
|