def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
conv_block = []
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim),
Mask(dim, mask_loss_type=self.opt.mask_loss_type),
nn.ReLU(True)]
if use_dropout:
conv_block += [nn.Dropout(0.5)]
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' % padding_type)
conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias),
norm_layer(dim),
Mask(dim, mask_loss_type=self.opt.mask_loss_type),
]
return nn.Sequential(*conv_block)
def __init__(self, input_nc=3, output_nc=3, ngf=64, norm_layer=nn.InstanceNorm2d, use_dropout=False, n_blocks=9, padding_type='reflect', opt=None):
super(MaskResnetGenerator, self).__init__()
assert(n_blocks >= 0)
self.opt = opt
self.device = torch.device(f'cuda:{opt.gpu_ids[0]}') if len(opt.gpu_ids) > 0 else 'cpu'
norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
if self.opt.upconv_bound:
upconv_mask_loss_type = 'bound'
else:
upconv_mask_loss_type = opt.mask_loss_type
model = [nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),
norm_layer(ngf),
Mask(ngf, mask_loss_type=upconv_mask_loss_type),
nn.ReLU(True)]
n_downsampling = 2
for i in range(n_downsampling):
mult = 2 ** i
model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
norm_layer(ngf * mult * 2),
Mask(ngf * mult * 2, mask_loss_type=self.opt.mask_loss_type if i == 1 else upconv_mask_loss_type),
nn.ReLU(True)]
mult = 2 ** n_downsampling
for i in range(n_blocks):
model += [MaskResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer,
use_dropout=use_dropout, use_bias=use_bias, opt=self.opt)]
for i in range(n_downsampling):
mult = 2 ** (n_downsampling - i)
if i == n_downsampling - 1 and self.opt.unmask_last_upconv:
model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
kernel_size=3, stride=2,
padding=1, output_padding=1,
bias=use_bias),
norm_layer(int(ngf * mult / 2)),
nn.ReLU(True)]
else:
model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),
kernel_size=3, stride=2,
padding=1, output_padding=1,
bias=use_bias),
norm_layer(int(ngf * mult / 2)),
Mask(int(ngf * mult / 2), mask_loss_type=upconv_mask_loss_type if i != n_downsampling-1 else 'bound'),
nn.ReLU(True)]
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
self.block_sparsity_coeff = torch.FloatTensor([1.0] * (self.opt.ngf * 4)).to(self.device)
def __init__(self, outer_nc, inner_nc, input_nc=None, submodule=None,
outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False, opt=None):
super(MaskUnetSkipConnectionBlock, self).__init__()
self.opt = opt
self.outermost = outermost
norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
if type(norm_layer) == functools.partial:
use_bias = norm_layer.func == nn.InstanceNorm2d
else:
use_bias = norm_layer == nn.InstanceNorm2d
if input_nc is None:
input_nc = outer_nc
downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
stride=2, padding=1, bias=use_bias)
downrelu = nn.LeakyReLU(0.2, True)
downnorm = norm_layer(inner_nc)
downmask = Mask(inner_nc, mask_loss_type=self.opt.mask_loss_type)
uprelu = nn.ReLU(True)
upnorm = norm_layer(outer_nc)
upmask = Mask(outer_nc, mask_loss_type=self.opt.mask_loss_type)
if outermost:
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
kernel_size=4, stride=2,
padding=1)
down = [downconv, downmask]
up = [uprelu, upconv, nn.Tanh()]
model = down + [submodule] + up
elif innermost:
upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
kernel_size=4, stride=2,
padding=1, bias=use_bias)
down = [downrelu, downconv, downmask]
up = [uprelu, upconv, upnorm, upmask]
model = down + up
else:
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
kernel_size=4, stride=2,
padding=1, bias=use_bias)
down = [downrelu, downconv, downnorm, downmask]
up = [uprelu, upconv, upnorm, upmask]
if use_dropout:
model = down + [submodule] + up + [nn.Dropout(0.5)]
else:
model = down + [submodule] + up
self.model = nn.Sequential(*model)
def build_conv_block(self, dim, padding_type, norm_layer, dropout_rate,
use_bias):
conv_block = []
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' %
padding_type)
conv_block += [
SeparableConv2d(in_channels=dim,
out_channels=dim,
kernel_size=3,
padding=p,
stride=1),
norm_layer(dim),
Mask(dim, mask_loss_type=self.opt.mask_loss_type),
nn.ReLU(True)
]
conv_block += [nn.Dropout(dropout_rate)]
p = 0
if padding_type == 'reflect':
conv_block += [nn.ReflectionPad2d(1)]
elif padding_type == 'replicate':
conv_block += [nn.ReplicationPad2d(1)]
elif padding_type == 'zero':
p = 1
else:
raise NotImplementedError('padding [%s] is not implemented' %
padding_type)
conv_block += [
SeparableConv2d(in_channels=dim,
out_channels=dim,
kernel_size=3,
padding=p,
stride=1),
norm_layer(dim),
Mask(dim, mask_loss_type=self.opt.mask_loss_type),
]
return nn.Sequential(*conv_block)