mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add negative_transforms switch to resgen

Browse Source
This commit is contained in:
James Betker 2020-06-23 09:41:12 -06:00
parent dfcbe5f2db
commit 0584c3b587
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
codes/models/archs/SwitchedResidualGenerator_arch.py
View File

@ -74,17 +74,23 @@ def create_sequential_growing_processing_block(filters_init, filter_growth, num_
class SwitchComputer(nn.Module):
def __init__(self, channels_in, filters, growth, transform_block, transform_count, reduction_blocks, processing_blocks=0, init_temp=20):
def __init__(self, channels_in, filters, growth, transform_block, transform_count, reduction_blocks, processing_blocks=0,
init_temp=20, enable_negative_transforms=False):
super(SwitchComputer, self).__init__()
self.enable_negative_transforms = enable_negative_transforms
self.filter_conv = ConvBnLelu(channels_in, filters)
self.reduction_blocks = nn.ModuleList([HalvingProcessingBlock(filters * 2 ** i) for i in range(reduction_blocks)])
final_filters = filters * 2 ** reduction_blocks
self.processing_blocks, final_filters = create_sequential_growing_processing_block(final_filters, growth, processing_blocks)
proc_block_filters = max(final_filters // 2, transform_count)
self.proc_switch_conv = ConvBnLelu(final_filters, proc_block_filters, bn=False)
self.final_switch_conv = nn.Conv2d(proc_block_filters, transform_count, 1, 1, 0)
tc = transform_count
if self.enable_negative_transforms:
tc = transform_count * 2
self.final_switch_conv = nn.Conv2d(proc_block_filters, tc, 1, 1, 0)
self.transforms = nn.ModuleList([transform_block() for i in range(transform_count)])
self.transforms = nn.ModuleList([transform_block() for _ in range(transform_count)])
# And the switch itself, including learned scalars
self.switch = BareConvSwitch(initial_temperature=init_temp)
@ -93,6 +99,8 @@ class SwitchComputer(nn.Module):
def forward(self, x, output_attention_weights=False):
xformed = [t.forward(x) for t in self.transforms]
if self.enable_negative_transforms:
xformed.extend([-t for t in xformed])
multiplexer = self.filter_conv(x)
for block in self.reduction_blocks: