def __init__(self, n_styles=18, opts=None):
super(ResNetProgressiveBackboneEncoder, self).__init__()
self.conv1 = nn.Conv2d(opts.input_nc,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = BatchNorm2d(64)
self.relu = PReLU(64)
resnet_basenet = resnet34(pretrained=True)
blocks = [
resnet_basenet.layer1, resnet_basenet.layer2,
resnet_basenet.layer3, resnet_basenet.layer4
]
modules = []
for block in blocks:
for bottleneck in block:
modules.append(bottleneck)
self.body = Sequential(*modules)
self.styles = nn.ModuleList()
self.style_count = n_styles
for i in range(self.style_count):
style = GradualStyleBlock(512, 512, 16)
self.styles.append(style)
self.progressive_stage = ProgressiveStage.Inference
def __init__(self, num_layers, mode='ir', n_styles=18, opts=None):
super(ProgressiveBackboneEncoder, self).__init__()
assert num_layers in [50, 100,
152], 'num_layers should be 50,100, or 152'
assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
blocks = get_blocks(num_layers)
if mode == 'ir':
unit_module = bottleneck_IR
elif mode == 'ir_se':
unit_module = bottleneck_IR_SE
self.input_layer = Sequential(
Conv2d(opts.input_nc, 64, (3, 3), 1, 1, bias=False),
BatchNorm2d(64), PReLU(64))
modules = []
for block in blocks:
for bottleneck in block:
modules.append(
unit_module(bottleneck.in_channel, bottleneck.depth,
bottleneck.stride))
self.body = Sequential(*modules)
self.styles = nn.ModuleList()
self.style_count = n_styles
for i in range(self.style_count):
style = GradualStyleBlock(512, 512, 16)
self.styles.append(style)
self.progressive_stage = ProgressiveStage.Inference
def __init__(self, n_styles=18, opts=None):
super(ResNetGradualStyleEncoder, self).__init__()
self.conv1 = nn.Conv2d(opts.input_nc,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = BatchNorm2d(64)
self.relu = PReLU(64)
resnet_basenet = resnet34(pretrained=True)
blocks = [
resnet_basenet.layer1, resnet_basenet.layer2,
resnet_basenet.layer3, resnet_basenet.layer4
]
modules = []
for block in blocks:
for bottleneck in block:
modules.append(bottleneck)
self.body = Sequential(*modules)
self.styles = nn.ModuleList()
self.style_count = n_styles
self.coarse_ind = 3
self.middle_ind = 7
for i in range(self.style_count):
if i < self.coarse_ind:
style = GradualStyleBlock(512, 512, 16)
elif i < self.middle_ind:
style = GradualStyleBlock(512, 512, 32)
else:
style = GradualStyleBlock(512, 512, 64)
self.styles.append(style)
self.latlayer1 = nn.Conv2d(256,
512,
kernel_size=1,
stride=1,
padding=0)
self.latlayer2 = nn.Conv2d(128,
512,
kernel_size=1,
stride=1,
padding=0)
def __init__(self, num_layers, mode='ir', n_styles=18, opts=None):
super(GradualStyleEncoder, self).__init__()
assert num_layers in [50, 100,
152], 'num_layers should be 50,100, or 152'
assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
blocks = get_blocks(num_layers)
if mode == 'ir':
unit_module = bottleneck_IR
elif mode == 'ir_se':
unit_module = bottleneck_IR_SE
self.input_layer = Sequential(
Conv2d(opts.input_nc, 64, (3, 3), 1, 1, bias=False),
BatchNorm2d(64), PReLU(64))
modules = []
for block in blocks:
for bottleneck in block:
modules.append(
unit_module(bottleneck.in_channel, bottleneck.depth,
bottleneck.stride))
self.body = Sequential(*modules)
self.styles = nn.ModuleList()
self.style_count = n_styles
self.coarse_ind = 3
self.middle_ind = 7
for i in range(self.style_count):
if i < self.coarse_ind:
style = GradualStyleBlock(512, 512, 16)
elif i < self.middle_ind:
style = GradualStyleBlock(512, 512, 32)
else:
style = GradualStyleBlock(512, 512, 64)
self.styles.append(style)
self.latlayer1 = nn.Conv2d(256,
512,
kernel_size=1,
stride=1,
padding=0)
self.latlayer2 = nn.Conv2d(128,
512,
kernel_size=1,
stride=1,
padding=0)