File size: 2,290 Bytes
61522a1 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import models.utils as mutils
from models import register
class Block(nn.Module):
def __init__(self, nf, group=1):
super(Block, self).__init__()
self.b1 = mutils.EResidualBlock(nf, nf, group=group)
self.c1 = mutils.BasicBlock(nf*2, nf, 1, 1, 0)
self.c2 = mutils.BasicBlock(nf*3, nf, 1, 1, 0)
self.c3 = mutils.BasicBlock(nf*4, nf, 1, 1, 0)
def forward(self, x):
c0 = o0 = x
b1 = self.b1(o0)
c1 = torch.cat([c0, b1], dim=1)
o1 = self.c1(c1)
b2 = self.b1(o1)
c2 = torch.cat([c1, b2], dim=1)
o2 = self.c2(c2)
b3 = self.b1(o2)
c3 = torch.cat([c2, b3], dim=1)
o3 = self.c3(c3)
return o3
@register('carn')
class CARN_M(nn.Module):
def __init__(self, in_nc=3, out_nc=3, nf=64, scale=4, group=4, no_upsampling=False):
super(CARN_M, self).__init__()
self.scale = scale
self.out_dim = nf
self.entry = nn.Conv2d(in_nc, nf, 3, 1, 1)
self.b1 = Block(nf, group=group)
self.b2 = Block(nf, group=group)
self.b3 = Block(nf, group=group)
self.c1 = mutils.BasicBlock(nf*2, nf, 1, 1, 0)
self.c2 = mutils.BasicBlock(nf*3, nf, 1, 1, 0)
self.c3 = mutils.BasicBlock(nf*4, nf, 1, 1, 0)
self.no_upsampling = no_upsampling
if not no_upsampling:
self.upsample = mutils.UpsampleBlock(nf, scale=scale, multi_scale=False, group=group)
self.exit = nn.Conv2d(nf, out_nc, 3, 1, 1)
def forward(self, x):
#x = self.sub_mean(x)
x = self.entry(x)
c0 = o0 = x
b1 = self.b1(o0)
c1 = torch.cat([c0, b1], dim=1)
o1 = self.c1(c1)
b2 = self.b2(o1)
c2 = torch.cat([c1, b2], dim=1)
o2 = self.c2(c2)
b3 = self.b3(o2)
c3 = torch.cat([c2, b3], dim=1)
o3 = self.c3(c3)
out = o3.clone()
if not self.no_upsampling:
out = self.upsample(out, scale=self.scale)
out = self.exit(out)
#out = self.add_mean(out)
return out
|