SRG2classic further re-integration
This commit is contained in:
James Betker
parent
9e2c96ad5d
commit
4e5ba61ae7
|
|
@ -2,201 +2,18 @@ import os
|
|||
|
||||
import torch
|
||||
import torchvision
|
||||
from matplotlib import cm
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
import functools
|
||||
from collections import OrderedDict
|
||||
|
||||
from torch.nn import init
|
||||
|
||||
from models.archs.arch_util import ConvBnLelu, ConvGnSilu
|
||||
from models.archs.SwitchedResidualGenerator_arch import HalvingProcessingBlock, ConfigurableSwitchComputer
|
||||
from models.archs.arch_util import ConvBnLelu, ConvGnSilu, ExpansionBlock, MultiConvBlock
|
||||
from switched_conv.switched_conv import BareConvSwitch, AttentionNorm
|
||||
from utils.util import checkpoint
|
||||
|
||||
|
||||
def initialize_weights(net_l, scale=1):
|
||||
if not isinstance(net_l, list):
|
||||
net_l = [net_l]
|
||||
for net in net_l:
|
||||
for m in net.modules():
|
||||
if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv3d):
|
||||
init.kaiming_normal_(m.weight, a=0, mode='fan_in')
|
||||
m.weight.data *= scale # for residual block
|
||||
if m.bias is not None:
|
||||
m.bias.data.zero_()
|
||||
elif isinstance(m, nn.Linear):
|
||||
init.kaiming_normal_(m.weight, a=0, mode='fan_in')
|
||||
m.weight.data *= scale
|
||||
if m.bias is not None:
|
||||
m.bias.data.zero_()
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
init.constant_(m.weight, 1)
|
||||
init.constant_(m.bias.data, 0.0)
|
||||
|
||||
|
||||
class AttentionNorm(nn.Module):
|
||||
def __init__(self, group_size, accumulator_size=128):
|
||||
super(AttentionNorm, self).__init__()
|
||||
self.accumulator_desired_size = accumulator_size
|
||||
self.group_size = group_size
|
||||
# These are all tensors so that they get saved with the graph.
|
||||
self.accumulator = nn.Parameter(torch.zeros(accumulator_size, group_size), requires_grad=False)
|
||||
self.accumulator_index = nn.Parameter(torch.zeros(1, dtype=torch.long, device='cpu'), requires_grad=False)
|
||||
self.accumulator_filled = nn.Parameter(torch.zeros(1, dtype=torch.bool, device='cpu'), requires_grad=False)
|
||||
|
||||
# Returns tensor of shape (group,) with a normalized mean across the accumulator in the range [0,1]. The intent
|
||||
# is to divide your inputs by this value.
|
||||
def compute_buffer_norm(self):
|
||||
if self.accumulator_filled:
|
||||
return torch.mean(self.accumulator, dim=0)
|
||||
else:
|
||||
return torch.ones(self.group_size, device=self.accumulator.device)
|
||||
|
||||
def add_norm_to_buffer(self, x):
|
||||
flat = x.sum(dim=[0, 1, 2], keepdim=True)
|
||||
norm = flat / torch.mean(flat)
|
||||
|
||||
# This often gets reset in GAN mode. We *never* want gradient accumulation in this parameter.
|
||||
self.accumulator.requires_grad = False
|
||||
self.accumulator[self.accumulator_index] = norm.detach()
|
||||
self.accumulator_index += 1
|
||||
if self.accumulator_index >= self.accumulator_desired_size:
|
||||
self.accumulator_index *= 0
|
||||
self.accumulator_filled |= True
|
||||
|
||||
# Input into forward is an attention tensor of shape (batch,width,height,groups)
|
||||
def forward(self, x: torch.Tensor):
|
||||
assert len(x.shape) == 4
|
||||
# Push the accumulator to the right device on the first iteration.
|
||||
if self.accumulator.device != x.device:
|
||||
self.accumulator = self.accumulator.to(x.device)
|
||||
|
||||
self.add_norm_to_buffer(x)
|
||||
norm = self.compute_buffer_norm()
|
||||
x = x / norm
|
||||
|
||||
# Need to re-normalize x so that the groups dimension sum to 1, just like when it was fed in.
|
||||
groups_sum = x.sum(dim=3, keepdim=True)
|
||||
return x / groups_sum
|
||||
|
||||
|
||||
class BareConvSwitch(nn.Module):
|
||||
"""
|
||||
Initializes the ConvSwitch.
|
||||
initial_temperature: The initial softmax temperature of the attention mechanism. For training from scratch, this
|
||||
should be set to a high number, for example 30.
|
||||
attention_norm: If specified, the AttentionNorm layer applied immediately after Softmax.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
initial_temperature=1,
|
||||
attention_norm=None
|
||||
):
|
||||
super(BareConvSwitch, self).__init__()
|
||||
|
||||
self.softmax = nn.Softmax(dim=-1)
|
||||
self.temperature = initial_temperature
|
||||
self.attention_norm = attention_norm
|
||||
|
||||
initialize_weights(self)
|
||||
|
||||
def set_attention_temperature(self, temp):
|
||||
self.temperature = temp
|
||||
|
||||
# SwitchedConv.forward takes these arguments;
|
||||
# conv_group: List of inputs (len=n) to the switch, each with shape (b,f,w,h)
|
||||
# conv_attention: Attention computation as an output from a conv layer, of shape (b,n,w,h). Before softmax
|
||||
# output_attention_weights: If True, post-softmax attention weights are returned.
|
||||
def forward(self, conv_group, conv_attention, output_attention_weights=False):
|
||||
# Stack up the conv_group input first and permute it to (batch, width, height, filter, groups)
|
||||
conv_outputs = torch.stack(conv_group, dim=0).permute(1, 3, 4, 2, 0)
|
||||
|
||||
conv_attention = conv_attention.permute(0, 2, 3, 1)
|
||||
conv_attention = self.softmax(conv_attention / self.temperature)
|
||||
if self.attention_norm:
|
||||
conv_attention = self.attention_norm(conv_attention)
|
||||
|
||||
# conv_outputs shape: (batch, width, height, filters, groups)
|
||||
# conv_attention shape: (batch, width, height, groups)
|
||||
# We want to format them so that we can matmul them together to produce:
|
||||
# desired shape: (batch, width, height, filters)
|
||||
# Note: conv_attention will generally be cast to float32 regardless of the input type, so cast conv_outputs to
|
||||
# float32 as well to match it.
|
||||
if self.training:
|
||||
# Doing it all in one op is substantially faster - better for training.
|
||||
attention_result = torch.einsum(
|
||||
"...ij,...j->...i", [conv_outputs.float(), conv_attention]
|
||||
)
|
||||
else:
|
||||
# eval_mode substantially reduces the GPU memory required to compute the attention result by performing the
|
||||
# attention multiplications one at a time. This is probably necessary for large images and attention breadths.
|
||||
attention_result = conv_outputs[:, :, :, :, 0] * conv_attention[:, :, :, 0].unsqueeze(dim=-1)
|
||||
for i in range(1, conv_attention.shape[-1]):
|
||||
attention_result += conv_outputs[:, :, :, :, i] * conv_attention[:, :, :, i].unsqueeze(dim=-1)
|
||||
|
||||
# Remember to shift the filters back into the expected slot.
|
||||
if output_attention_weights:
|
||||
return attention_result.permute(0, 3, 1, 2), conv_attention
|
||||
else:
|
||||
return attention_result.permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
class MultiConvBlock(nn.Module):
|
||||
def __init__(self, filters_in, filters_mid, filters_out, kernel_size, depth, scale_init=1.0, norm=False, weight_init_factor=1):
|
||||
assert depth >= 2
|
||||
super(MultiConvBlock, self).__init__()
|
||||
self.noise_scale = nn.Parameter(torch.full((1,), fill_value=.01))
|
||||
self.bnconvs = nn.ModuleList([ConvBnLelu(filters_in, filters_mid, kernel_size, norm=norm, bias=False, weight_init_factor=weight_init_factor)] +
|
||||
[ConvBnLelu(filters_mid, filters_mid, kernel_size, norm=norm, bias=False, weight_init_factor=weight_init_factor) for i in range(depth - 2)] +
|
||||
[ConvBnLelu(filters_mid, filters_out, kernel_size, activation=False, norm=False, bias=False, weight_init_factor=weight_init_factor)])
|
||||
self.scale = nn.Parameter(torch.full((1,), fill_value=scale_init))
|
||||
self.bias = nn.Parameter(torch.zeros(1), requires_grad=False)
|
||||
|
||||
def forward(self, x, noise=None):
|
||||
if noise is not None:
|
||||
noise = noise * self.noise_scale
|
||||
x = x + noise
|
||||
for m in self.bnconvs:
|
||||
x = m.forward(x)
|
||||
return x * self.scale + self.bias
|
||||
|
||||
|
||||
# VGG-style layer with Conv(stride2)->BN->Activation->Conv->BN->Activation
|
||||
# Doubles the input filter count.
|
||||
class HalvingProcessingBlock(nn.Module):
|
||||
def __init__(self, filters):
|
||||
super(HalvingProcessingBlock, self).__init__()
|
||||
self.bnconv1 = ConvGnSilu(filters, filters * 2, stride=2, norm=False, bias=False)
|
||||
self.bnconv2 = ConvGnSilu(filters * 2, filters * 2, norm=True, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.bnconv1(x)
|
||||
return self.bnconv2(x)
|
||||
|
||||
|
||||
# Block that upsamples 2x and reduces incoming filters by 2x. It preserves structure by taking a passthrough feed
|
||||
# along with the feature representation.
|
||||
class ExpansionBlock(nn.Module):
|
||||
def __init__(self, filters_in, filters_out=None, block=ConvGnSilu):
|
||||
super(ExpansionBlock, self).__init__()
|
||||
if filters_out is None:
|
||||
filters_out = filters_in // 2
|
||||
self.decimate = block(filters_in, filters_out, kernel_size=1, bias=False, activation=False, norm=True)
|
||||
self.process_passthrough = block(filters_out, filters_out, kernel_size=3, bias=True, activation=False, norm=True)
|
||||
self.conjoin = block(filters_out*2, filters_out, kernel_size=3, bias=False, activation=True, norm=False)
|
||||
self.process = block(filters_out, filters_out, kernel_size=3, bias=False, activation=True, norm=True)
|
||||
|
||||
# input is the feature signal with shape (b, f, w, h)
|
||||
# passthrough is the structure signal with shape (b, f/2, w*2, h*2)
|
||||
# output is conjoined upsample with shape (b, f/2, w*2, h*2)
|
||||
def forward(self, input, passthrough):
|
||||
x = F.interpolate(input, scale_factor=2, mode="nearest")
|
||||
x = self.decimate(x)
|
||||
p = self.process_passthrough(passthrough)
|
||||
x = self.conjoin(torch.cat([x, p], dim=1))
|
||||
return self.process(x)
|
||||
|
||||
|
||||
# This is a classic u-net architecture with the goal of assigning each individual pixel an individual transform
|
||||
# switching set.
|
||||
class ConvBasisMultiplexer(nn.Module):
|
||||
|
|
@ -231,50 +48,6 @@ class ConvBasisMultiplexer(nn.Module):
|
|||
return x
|
||||
|
||||
|
||||
class ConfigurableSwitchComputer(nn.Module):
|
||||
def __init__(self, base_filters, multiplexer_net, pre_transform_block, transform_block, transform_count, init_temp=20,
|
||||
add_scalable_noise_to_transforms=False):
|
||||
super(ConfigurableSwitchComputer, self).__init__()
|
||||
|
||||
tc = transform_count
|
||||
self.multiplexer = multiplexer_net(tc)
|
||||
|
||||
self.pre_transform = pre_transform_block()
|
||||
self.transforms = nn.ModuleList([transform_block() for _ in range(transform_count)])
|
||||
self.add_noise = add_scalable_noise_to_transforms
|
||||
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, attention_norm=AttentionNorm(transform_count, accumulator_size=16 * transform_count))
|
||||
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)
|
||||
# depending on its needs.
|
||||
self.psc_scale = nn.Parameter(torch.full((1,), float(.1)))
|
||||
|
||||
def forward(self, x, output_attention_weights=True):
|
||||
identity = x
|
||||
if self.add_noise:
|
||||
rand_feature = torch.randn_like(x) * self.noise_scale
|
||||
x = x + rand_feature
|
||||
|
||||
x = self.pre_transform(x)
|
||||
xformed = [t.forward(x) for t in self.transforms]
|
||||
|
||||
m = self.multiplexer(identity)
|
||||
|
||||
outputs, attention = self.switch(xformed, m, True)
|
||||
outputs = identity + outputs * self.switch_scale
|
||||
outputs = outputs + self.post_switch_conv(outputs) * self.psc_scale
|
||||
if output_attention_weights:
|
||||
return outputs, attention
|
||||
else:
|
||||
return outputs
|
||||
|
||||
def set_temperature(self, temp):
|
||||
self.switch.set_attention_temperature(temp)
|
||||
|
||||
|
||||
def compute_attention_specificity(att_weights, topk=3):
|
||||
att = att_weights.detach()
|
||||
vals, indices = torch.topk(att, topk, dim=-1)
|
||||
|
|
@ -347,7 +120,7 @@ class ConfigurableSwitchedResidualGenerator2(nn.Module):
|
|||
multiplx_fn = functools.partial(ConvBasisMultiplexer, transformation_filters, switch_filters, switch_reductions, switch_processing_layers, trans_counts)
|
||||
pretransform_fn = functools.partial(ConvBnLelu, transformation_filters, transformation_filters, norm=False, bias=False, weight_init_factor=.1)
|
||||
transform_fn = functools.partial(MultiConvBlock, transformation_filters, int(transformation_filters * 1.5), transformation_filters, kernel_size=trans_kernel_sizes, depth=trans_layers, weight_init_factor=.1)
|
||||
switches.append(ConfigurableSwitchComputer(transformation_filters, multiplx_fn,
|
||||
switches.append(ConfigurableSwitchComputer(transformation_filters, multiplx_fn, attention_norm=True,
|
||||
pre_transform_block=pretransform_fn, transform_block=transform_fn,
|
||||
transform_count=trans_counts, init_temp=initial_temp,
|
||||
add_scalable_noise_to_transforms=add_scalable_noise_to_transforms))
|
||||
|
|
@ -412,10 +185,11 @@ class ConfigurableSwitchedResidualGenerator2(nn.Module):
|
|||
|
||||
|
||||
class Interpolate(nn.Module):
|
||||
def __init__(self, factor):
|
||||
def __init__(self, factor, mode="nearest"):
|
||||
super(Interpolate, self).__init__()
|
||||
self.factor = factor
|
||||
self.mode = mode
|
||||
|
||||
def forward(self, x):
|
||||
return F.interpolate(x, scale_factor=self.factor)
|
||||
return F.interpolate(x, scale_factor=self.factor, mode=self.mode)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user