| """Blue Marble model architecture (Apache-2.0).""" |
| import torch, torch.nn as nn, torch.nn.functional as F |
| S2_BANDS = ['B01', 'B02', 'B03', 'B04', 'B05', 'B06', 'B07', 'B08', 'B8A', 'B09', 'B10', 'B11', 'B12'] |
|
|
| def conv_bn(i, o, k=3, s=1, g=1, d=1, act=True): |
| layers = [nn.Conv2d(i, o, k, s, d * (k // 2), dilation=d, groups=g, bias=False), nn.BatchNorm2d(o)] |
| if act: |
| layers.append(nn.ReLU6(inplace=True)) |
| return nn.Sequential(*layers) |
|
|
| class SE(nn.Module): |
|
|
| def __init__(self, c, r=4): |
| super().__init__() |
| h = max(8, c // r) |
| self.f = nn.Sequential(nn.Conv2d(c, h, 1), nn.ReLU6(inplace=True), nn.Conv2d(h, c, 1), nn.Hardsigmoid()) |
|
|
| def forward(self, x): |
| return x * self.f(x.mean((2, 3), keepdim=True)) |
|
|
| class IRB(nn.Module): |
|
|
| def __init__(self, i, o, s=1, e=3, k=5, se=False, d=1): |
| super().__init__() |
| h = int(round(i * e)) |
| self.res = s == 1 and i == o |
| self.net = nn.Sequential(conv_bn(i, h, 1) if e != 1 else nn.Identity(), conv_bn(h, h, k, s, g=h, d=d), SE(h) if se else nn.Identity(), conv_bn(h, o, 1, act=False)) |
|
|
| def forward(self, x): |
| y = self.net(x) |
| return x + y if self.res else y |
|
|
| class SpectralStem(nn.Module): |
|
|
| def __init__(self, c, in_bands=13): |
| super().__init__() |
| self.spectral = nn.Conv2d(in_bands, c, 1, bias=False) |
| self.bn = nn.BatchNorm2d(c) |
| self.act = nn.ReLU6(inplace=True) |
| self.spatial = conv_bn(c, c, 3, g=c) |
| self.fuse = conv_bn(c, c, 1) |
|
|
| def forward(self, x): |
| return self.fuse(self.spatial(self.act(self.bn(self.spectral(x))))) |
|
|
| @torch.no_grad() |
| def slice_bands(self, idx): |
| new = nn.Conv2d(len(idx), self.spectral.out_channels, 1, bias=False).to(self.spectral.weight.device) |
| new.weight.copy_(self.spectral.weight[:, idx]) |
| self.spectral = new |
|
|
| class Attn2d(nn.Module): |
|
|
| def __init__(self, c, heads=4, mlp=2): |
| super().__init__() |
| self.h, self.d = (heads, c // heads) |
| self.n1, self.q, self.k, self.v, self.o = (nn.BatchNorm2d(c), nn.Conv2d(c, c, 1), nn.Conv2d(c, c, 1), nn.Conv2d(c, c, 1), nn.Conv2d(c, c, 1)) |
| self.mlp = nn.Sequential(nn.BatchNorm2d(c), nn.Conv2d(c, c * mlp, 1), nn.ReLU6(inplace=True), nn.Conv2d(c * mlp, c, 1)) |
|
|
| def forward(self, x): |
| b, c, h, w = x.shape |
| n = self.n1(x) |
| sp = lambda t: t.reshape(b, self.h, self.d, h * w).transpose(2, 3) |
| q, k, v = (sp(self.q(n)), sp(self.k(n)), sp(self.v(n))) |
| a = (q.float() @ k.float().transpose(2, 3) / self.d ** 0.5).softmax(-1) |
| y = (a @ v.float()).to(x.dtype).transpose(2, 3).reshape(b, c, h, w) |
| x = x + self.o(y) |
| return x + self.mlp(x) |
| ARCHS = {'pico': dict(stem=8, w=(8, 16, 24, 40), d=(1, 2, 2, 2), e=2, k=3, rates=(1,), dec=16, full=False, se=False), 'nano': dict(stem=12, w=(16, 24, 48, 80), d=(1, 2, 3, 2), e=3, k=5, rates=(1, 2), dec=24, full=True, se=False), 'tiny': dict(stem=16, w=(24, 40, 80, 160), d=(2, 2, 4, 3), e=3, k=5, rates=(1, 2, 4), dec=40, full=True, se=True), 'small': dict(stem=24, w=(32, 56, 112, 208), d=(2, 3, 5, 3), e=3, k=7, rates=(1, 2, 4), dec=48, full=True, se=True), 'medium': dict(stem=32, w=(48, 96, 192, 320), d=(2, 3, 8, 4), e=4, k=7, rates=(1, 2, 4), dec=96, full=True, se=True, attn=2)} |
|
|
| class Encoder(nn.Module): |
|
|
| def __init__(self, cfg, in_bands=13): |
| super().__init__() |
| self.stem = SpectralStem(cfg['stem'], in_bands) |
| prev, self.stages = (cfg['stem'], nn.ModuleList()) |
| for si, (w, n) in enumerate(zip(cfg['w'], cfg['d'])): |
| blocks = [IRB(prev, w, 2, cfg['e'], 3)] |
| blocks += [IRB(w, w, 1, cfg['e'], cfg['k'], se=cfg['se'] and si >= 2) for _ in range(n - 1)] |
| if si == len(cfg['w']) - 1: |
| blocks += [Attn2d(w, heads=8) for _ in range(cfg.get('attn', 0))] |
| self.stages.append(nn.Sequential(*blocks)) |
| prev = w |
| self.neck = nn.ModuleList((IRB(prev, prev, 1, 2, 3, d=r) for r in cfg['rates'])) |
| self.channels = (cfg['stem'], *cfg['w']) |
|
|
| def forward(self, x): |
| out = [self.stem(x)] |
| for s in self.stages: |
| out.append(s(out[-1])) |
| y = out[-1] |
| for n in self.neck: |
| y = n(y) |
| out[-1] = y |
| return out |
|
|
| class FPNDecoder(nn.Module): |
|
|
| def __init__(self, chs, dec, full=True): |
| super().__init__() |
| self.full = full |
| self.lat = nn.ModuleList((conv_bn(c, dec, 1) for c in chs[1:])) |
| self.ref = nn.ModuleList((IRB(dec, dec, 1, 2, 3) for _ in chs[1:-1])) |
| if full: |
| self.lat0 = conv_bn(chs[0], dec, 1) |
| self.ref0 = IRB(dec, dec, 1, 2, 3) |
| self.out_channels = dec |
|
|
| def forward(self, f): |
| y = self.lat[-1](f[-1]) |
| for i in range(len(self.ref) - 1, -1, -1): |
| y = F.interpolate(y, scale_factor=2.0, mode='bilinear', align_corners=False) + self.lat[i](f[i + 1]) |
| y = self.ref[i](y) |
| if self.full: |
| y = F.interpolate(y, scale_factor=2.0, mode='bilinear', align_corners=False) + self.lat0(f[0]) |
| y = self.ref0(y) |
| return y |
|
|
| class BlueMarble(nn.Module): |
|
|
| def __init__(self, arch='tiny', num_classes=2, in_bands=13): |
| super().__init__() |
| cfg = ARCHS[arch] |
| self.arch, self.upsample = (arch, 1 if cfg['full'] else 2) |
| self.encoder = Encoder(cfg, in_bands) |
| self.decoder = FPNDecoder(self.encoder.channels, cfg['dec'], cfg['full']) |
| self.head = nn.Sequential(nn.Dropout2d(0.1), nn.Conv2d(cfg['dec'], num_classes, 1)) |
|
|
| def forward(self, x): |
| y = self.head(self.decoder(self.encoder(x))) |
| if self.upsample > 1: |
| y = F.interpolate(y, scale_factor=float(self.upsample), mode='bilinear', align_corners=False) |
| return y |
|
|
| def fast_layout(model, in_shape, device='cuda'): |
| import torch.nn.functional as F |
| if not str(device).startswith('cuda'): |
| return (model, False) |
| try: |
| m = model.to(memory_format=torch.channels_last) |
| x = torch.randn(in_shape, device=device).to(memory_format=torch.channels_last) |
| with torch.autocast('cuda', dtype=torch.bfloat16): |
| out = m(x) |
| out.float().mean().backward() |
| m.zero_grad(set_to_none=True) |
| return (m, True) |
| except RuntimeError as e: |
| print('channels_last unavailable, using NCHW:', str(e)[:120], flush=True) |
| return (model.to(memory_format=torch.contiguous_format), False) |
|
|
