def __init__(self, in_c, k=4, convGamma=True):
super(Attention_concat, self).__init__()
self.in_c = in_c
self.inter_c = in_c // k
self.query_conv = nn.Conv2d(in_channels=self.in_c,
out_channels=self.in_c // k,
kernel_size=1)
self.key_conv = nn.Conv2d(in_channels=self.in_c,
out_channels=self.in_c // k,
kernel_size=1)
self.concat_conv = nn.Conv2d(in_channels=self.inter_c * 2,
out_channels=1,
kernel_size=1)
self.value_conv = spectral_norm_fc(nn.Conv2d(in_channels=in_c,
out_channels=in_c,
kernel_size=1),
coeff=.9,
n_power_iterations=5)
self.convGamma = convGamma
if convGamma:
self.gamma = spectral_norm_fc(nn.Conv2d(in_channels=in_c,
out_channels=in_c,
kernel_size=1),
coeff=.9,
n_power_iterations=5)
else:
self.gamma = Parameter(torch.zeros(1))
self.nonlin_2 = nn.Tanh()
self.nonlin_1 = nn.ELU()
def __init__(self, input_channel_num, k=4, convGamma=True):
super(Attention_dot, self).__init__()
self.c_in = input_channel_num
self.query_conv = spectral_norm_fc(nn.Conv2d(in_channels=self.c_in,
out_channels=self.c_in //
k,
kernel_size=1),
coeff=.9,
n_power_iterations=5)
self.key_conv = spectral_norm_fc(nn.Conv2d(in_channels=self.c_in,
out_channels=self.c_in // k,
kernel_size=1),
coeff=.9,
n_power_iterations=5)
self.value_conv = spectral_norm_fc(nn.Conv2d(in_channels=self.c_in,
out_channels=self.c_in,
kernel_size=1),
coeff=.5,
n_power_iterations=5)
self.nonlin = nn.ELU()
self.convGamma = convGamma
if convGamma:
self.gamma = spectral_norm_fc(nn.Conv2d(in_channels=self.c_in,
out_channels=self.c_in,
kernel_size=1),
coeff=.9,
n_power_iterations=5)
else:
self.gamma = Parameter(torch.zeros(1))
self.nonlin_2 = nn.Tanh()
def _wrapper_spectral_norm(self, layer, shapes, kernel_size):
if kernel_size == 1:
# use spectral norm fc, because bound are tight for 1x1 convolutions
return spectral_norm_fc(layer, self.coeff,
n_power_iterations=self.n_power_iter)
else:
# use spectral norm based on conv, because bound not tight
return spectral_norm_conv(layer, self.coeff, shapes,
n_power_iterations=self.n_power_iter)
def _spec_norm_wrapper(self, layer):
return spectral_norm_fc(layer, coeff=.9, n_power_iterations=5)
