From dd0bbd9a7ca70d5502bddc6a0fa9392b58f77d22 Mon Sep 17 00:00:00 2001 From: James Betker Date: Mon, 13 Jul 2020 08:38:17 -0600 Subject: [PATCH] Enable AttentionNorm on SRG2 --- codes/models/archs/SwitchedResidualGenerator_arch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/models/archs/SwitchedResidualGenerator_arch.py b/codes/models/archs/SwitchedResidualGenerator_arch.py index c214819d..ea0a167c 100644 --- a/codes/models/archs/SwitchedResidualGenerator_arch.py +++ b/codes/models/archs/SwitchedResidualGenerator_arch.py @@ -123,7 +123,7 @@ class ConfigurableSwitchComputer(nn.Module): self.noise_scale = nn.Parameter(torch.full((1,), float(1e-3))) # And the switch itself, including learned scalars - self.switch = BareConvSwitch(initial_temperature=init_temp) + self.switch = BareConvSwitch(initial_temperature=init_temp, attention_norm=AttentionNorm(transform_count, accumulator_size=128)) self.switch_scale = nn.Parameter(torch.full((1,), float(1))) self.post_switch_conv = ConvBnLelu(base_filters, base_filters, norm=False, bias=True) # The post_switch_conv gets a low scale initially. The network can decide to magnify it (or not)