def layer_conv3d(self, net, args, options):
if len(args) > 0:
channels = args[0]
else:
channels = self.current_size.channels
options = hc.Config(options)
stride = options.stride or 1
fltr = options.filter or 3
dilation = 1
padding = options.padding or 1 #self.get_same_padding(self.current_width, self.current_width, stride, dilation)
if options.padding0:
padding = [options.padding0, padding, padding]
if options.stride0:
stride = [options.stride0, stride, stride]
else:
stride = [stride, stride, stride]
layers = [
nn.Conv3d(options.input_channels or self.current_size.channels,
channels,
fltr,
stride,
padding=padding)
]
self.nn_init(layer, options.initializer)
self.current_size = LayerShape(
frames, channels, self.current_size.height // stride[1],
self.current_size.width // stride[2]
) #TODO this doesn't work, what is frames? Also chw calculation like conv2d
return nn.Sequential(*layers)
def layer_equal_linear(self, net, args, options):
lr_mul = 1
if options.lr_mul is not None:
lr_mul = options.lr_mul
result = EqualLinear(self.current_size.size(), args[0], lr_mul=lr_mul)
self.current_size = LayerShape(args[0])
return result
def layer_linear(self, net, args, options):
options = hc.Config(options)
shape = [int(x) for x in str(args[0]).split("*")]
bias = True
if options.bias == False:
bias = False
output_size = 1
for dim in shape:
output_size *= dim
layers = []
if len(self.current_size.dims) != 1:
layers += [nn.Flatten()]
layers += [
nn.Linear(options.input_size or self.current_size.size(),
output_size,
bias=bias)
]
self.nn_init(layers[-1], options.initializer)
self.current_size = LayerShape(*list(reversed(shape)))
if len(shape) != 1:
layers.append(Reshape(*self.current_size.dims))
if self.is_latent:
self._latent_parameters += [layers[0].weight]
self.is_latent = False
return nn.Sequential(*layers)
def layer_conv1d(self, net, args, options):
if len(args) > 0:
channels = args[0]
else:
channels = self.current_size.channels
print("Options:", options)
options = hc.Config(options)
stride = options.stride or 1
fltr = options.filter or 3
dilation = 1
padding = 1
if options.padding is not None:
padding = options.padding
layers = [
nn.Conv1d(options.input_channels or self.current_size.channels,
channels,
fltr,
stride,
padding=padding)
]
self.nn_init(layers[-1], options.initializer)
h, _ = self.conv_output_shape(
(self.current_size.height, self.current_size.height),
options.filter or 3, stride, padding, 1)
self.current_size = LayerShape(channels, h)
return nn.Sequential(*layers)
def layer_conv2d(self, net, args, options):
if len(args) > 0:
channels = args[0]
else:
channels = self.current_size.channels
options = hc.Config(options)
stride = 1
if options.stride is not None:
stride = options.stride
filter = 3
if options.filter is not None:
filter = options.filter
padding = 1
if options.padding is not None:
padding = options.padding
dilation = 1
layer = nn.Conv2d(options.input_channels or self.current_size.channels,
channels,
filter,
stride,
padding=(padding, padding))
self.nn_init(layer, options.initializer)
h, w = self.conv_output_shape(
(self.current_size.height, self.current_size.width), filter,
stride, padding, dilation)
self.current_size = LayerShape(channels, h, w)
return layer
def layer_adaptive_avg_pool3d(self, net, args, options):
frames = 4 #TODO
self.current_size = LayerShape(frames, self.current_size.channels,
self.current_size.height // 2,
self.current_size.width // 2)
return nn.AdaptiveAvgPool3d([
self.current_size.frames, self.current_size.height,
self.current_size.width
]) #TODO looks wrong
def __init__(self, component, args, options):
super(Upsample, self).__init__(component, args, options)
h = options.h or component.current_size.height * 2
if len(component.current_size.dims) == 3:
w = options.w or component.current_size.width * 2
self.layer = nn.Upsample((h, w), mode="bilinear")
self.size = LayerShape(component.current_size.channels, h, w)
elif len(component.current_size.dims) == 2:
self.layer = nn.Upsample((h), mode="linear")
self.size = LayerShape(component.current_size.channels, h)
def layer_subpixel(self, net, args, options):
options = hc.Config(options)
channels = args[0]
layers = [
nn.Conv2d(options.input_channels or self.current_size.channels,
channels * 4, options.filter or 3, 1, 1),
nn.PixelShuffle(2)
]
self.nn_init(layers[0], options.initializer)
self.current_size = LayerShape(channels, self.current_size.height * 2,
self.current_size.width * 2)
return nn.Sequential(*layers)
def forward(self, input, context={}):
if self.get_device().index != input.device.index:
input = input.to(self.get_device())
for module, parsed, layer_shape in zip(self.net, self.parsed_layers,
self.layer_shapes):
try:
options = parsed.parsed_options
args = parsed.args
layer_name = parsed.layer_name
name = options.name
if isinstance(module, hg.Layer):
input = module(input, context)
elif layer_name == "adaptive_instance_norm":
input = module(input, context['w'])
elif layer_name == "ez_norm":
input = module(input, context['w'])
elif layer_name == "split":
input = torch.split(input, args[0], options.dim
or -1)[args[1]]
elif layer_name == "latent":
input = self.gan.latent.z #sample()
elif layer_name == "modulated_conv2d":
input = module(input, context['w'])
elif layer_name == "pretrained":
in_zero_one = (input + self.const_one) / self.const_two
mean = torch.as_tensor([0.485, 0.456, 0.406],
device='cuda:0').view(1, 3, 1, 1)
std = torch.as_tensor([0.229, 0.224, 0.225],
device='cuda:0').view(1, 3, 1, 1)
input = module(input.clone().sub_(mean).div_(std))
else:
input = module(input)
if self.gan.steps == 0:
size = LayerShape(*list(input.shape[1:]))
if size.squeeze_dims() != layer_shape.squeeze_dims():
print("Error: Size error on", layer_name)
print("Error: Expected output size", layer_shape.dims)
print("Error: Actual output size", size.dims)
raise "Layer size error, cannot continue"
else:
pass
if name is not None:
context[name] = input
except:
raise ValidationException(
"Error on " + parsed.layer_defn + " - input size " +
",".join([str(x) for x in input.shape]))
self.sample = input
return input
def layer_resize_conv2d(self, net, args, options):
options = hc.Config(options)
channels = args[0]
w = options.w or self.current_size.width * 2
h = options.h or self.current_size.height * 2
layers = [
nn.Upsample((h, w), mode="bilinear"),
nn.Conv2d(options.input_channels or self.current_size.channels,
channels, options.filter or 3, 1, 1)
]
self.nn_init(layers[-1], options.initializer)
self.current_size = LayerShape(channels, h, w)
return nn.Sequential(*layers)
def layer_split(self, net, args, options):
options = hc.Config(options)
split_size = args[0]
select = args[1]
dim = -1
if options.dim:
dim = options.dim
#TODO better validation
#TODO increase dim options
if dim == -1:
dims = list(self.current_size.dims).copy()
dims[0] = split_size
if (select + 1) * split_size > self.current_size.channels:
dims[0] = self.current_size.channels % split_size
self.current_size = LayerShape(*dims)
return NoOp()
def __init__(self, component, args, options):
super(ResizableStack, self).__init__(component, args, options)
self.size = LayerShape(component.gan.channels(), component.gan.height(), component.gan.width())
self.max_channels = options.max_channels or 256
self.segment_channels = options.segment_channels or 5
self.style = options.style or "w"
layers = []
sizes = self.sizes(component.current_size.height, component.current_size.width, component.gan.height(), component.gan.width(), self.segment_channels * 2 * component.gan.channels())
print("SIZES", sizes)
for i, size in enumerate(sizes[1:]):
c = min(size.channels, self.max_channels)
upsample = hg.layers.Upsample(component, [], hc.Config({"w": size.width, "h": size.height}))
component.current_size = upsample.output_size() #TODO abstract input_size
if options.normalize != False:
_, add = component.parse_layer("add self (ez_norm initializer=(xavier_normal) style=" + self.style + ")")
_, conv = component.parse_layer("conv2d " + str(size.channels) + " padding=0 initializer=(xavier_normal)")
if options.normalize == False:
layers += [upsample, conv]
else:
layers += [upsample, add, conv]
if i < len(sizes) - 2:
layers += [nn.ReLU()]
layers += [hg.layers.SegmentSoftmax(component, [component.gan.channels()], {})]
self.layers = nn.ModuleList(layers)
def sizes(self, initial_height, initial_width, target_height, target_width, final_channels):
channels = []
hs = []
ws = []
sizes = []
w = initial_width
h = initial_height
i = 0
channels.append(final_channels)
while w < target_width or h < target_height:
if i > 0:
h-=2 #padding
w-=2 #padding
h*=2 #upscale
w*=2 #upscale
channels.append(final_channels * 2**i)
hs.append(min(h, target_height+2))
ws.append(min(w, target_width+2))
i+=1
w = initial_width
h = initial_height
channels.reverse()
for c,h,w in zip(channels, hs, ws):
sizes.append(LayerShape(c, h, w))
return sizes
def create(self):
self.latent = self.create_component("latent")
self.generator = self.create_component(
"generator",
input=self.latent,
context_shapes={"q": LayerShape(self.config.length // 2)})
self.discriminator = self.create_component("discriminator")
def __init__(self, component, args, options):
super(Layer, self).__init__(component, args, options)
self.name = args[0]
self.size = component.layer_output_sizes[args[0]]
if options.upsample:
self.size = LayerShape(self.size.channels, *component.current_size.dims[1:])
self.upsample = nn.Upsample(self.size.dims[1:], mode="bilinear")
def create(self):
self.generator = self.create_component("generator",
input=self.inputs.next()[0])
if self.config.generator2:
self.generator2 = self.create_component(
"generator2", input=self.inputs.next()[1])
self.discriminator = self.create_component(
"discriminator", context_shapes={"digit": LayerShape(10)})
self.loss = self.create_component("loss")
def layer_scaled_conv2d(self, net, args, options):
channels = self.current_size.channels
if len(args) > 0:
channels = args[0]
method = "conv"
if len(args) > 1:
method = args[1]
upsample = method == "upsample"
downsample = method == "downsample"
demodulate = True
if options.demodulate == False:
demodulate = False
filter = 3
if options.filter:
filter = options.filter
lr_mul = 1.0
if options.lr_mul:
lr_mul = options.lr_mul
input_channels = self.current_size.channels
if options.input_channels:
input_channels = options.input_channels
result = ScaledConv2d(input_channels,
channels,
filter,
0,
upsample=upsample,
demodulate=demodulate,
downsample=downsample,
lr_mul=lr_mul)
self.nn_init(result, options.initializer)
if upsample:
self.current_size = LayerShape(channels,
self.current_size.height * 2,
self.current_size.width * 2)
else:
self.current_size = LayerShape(channels,
self.current_size.height - 2,
self.current_size.width - 2)
return result
def layer_multi_head_attention(self, net, args, options):
output_size = self.current_size.size()
if len(args) > 0:
output_size = args[0]
layer = MultiHeadAttention(self.current_size.size(),
output_size,
heads=options.heads or 4)
self.current_size = LayerShape(output_size)
self.nn_init(layer.o, options.initializer)
self.nn_init(layer.h, options.initializer)
self.nn_init(layer.g, options.initializer)
self.nn_init(layer.f, options.initializer)
if self.is_latent:
self._latent_parameters += [
layer.h.weight, layer.g.weight, layer.f.weight
]
self.is_latent = False
return layer
def layer_modulated_conv2d(self, net, args, options):
channels = self.current_size.channels
if len(args) > 0:
channels = args[0]
method = "conv"
if len(args) > 1:
method = args[1]
upsample = method == "upsample"
downsample = method == "downsample"
demodulate = True
if options.demodulate == False:
demodulate = False
filter = 3
if options.filter:
filter = options.filter
lr_mul = 1.0
if options.lr_mul:
lr_mul = options.lr_mul
input_channels = self.current_size.channels
if options.input_channels:
input_channels = options.input_channels
result = ModulatedConv2d(input_channels,
channels,
filter,
self.layer_output_sizes['w'].size(),
upsample=upsample,
demodulate=demodulate,
downsample=downsample,
lr_mul=lr_mul)
if upsample:
self.current_size = LayerShape(channels,
self.current_size.height * 2,
self.current_size.width * 2)
elif downsample:
self.current_size = LayerShape(channels,
self.current_size.height // 2,
self.current_size.width // 2)
return result
def create(self):
self.latent = self.create_component("latent")
self.encoder = self.create_component("encoder")
self.decoder = self.create_component("decoder", input=self.encoder)
c_w = 8
c_h = 8
c_channels = 512
c_shape = LayerShape(c_channels, c_h, c_w)
self.state = self.create_component("state",
input=self.encoder,
context_shapes={"past": c_shape})
self.discriminator = self.create_component("discriminator")
def layer_deconv(self, net, args, options):
if len(args) > 0:
channels = args[0]
else:
channels = self.current_size.channels
options = hc.Config(options)
filter = 4 #TODO
if options.filter:
filter = options.filter
stride = 2
if options.stride:
stride = options.stride
padding = 1
if options.padding:
padding = options.padding
layer = nn.ConvTranspose2d(
options.input_channels or self.current_size.channels, channels,
filter, stride, padding)
self.nn_init(layer, options.initializer)
self.current_size = LayerShape(channels, self.current_size.height * 2,
self.current_size.width * 2)
return layer
def layer_resize_conv1d(self, net, args, options):
options = hc.Config(options)
channels = args[0]
h = options.h or self.current_size.height * 2
padding = 1
if options.padding is not None:
padding = options.padding
layers = [
nn.Upsample((h)),
nn.Conv1d(options.input_channels or self.current_size.channels,
channels,
options.filter or 3,
1,
padding=padding)
]
self.nn_init(layers[-1], options.initializer)
h, _ = self.conv_output_shape((h, h), options.filter or 3, 1, padding,
1)
self.current_size = LayerShape(channels, h)
return nn.Sequential(*layers)
def __init__(self, component, args, options):
super(EzNorm, self).__init__(component, args, options)
self.dim = options.dim or 1
self.size = LayerShape(*component.current_size.dims)
style_size = component.layer_output_sizes[options.style or 'w'].size()
channels = component.current_size.channels
dims = len(component.current_size.dims)
self.beta = nn.Linear(style_size, channels)
if dims == 2:
self.conv = nn.Conv1d(channels, 1, 1, 1, padding=0)
else:
self.conv = nn.Conv2d(channels, 1, 1, 1, padding=0)
component.nn_init(self.beta, options.initializer)
component.nn_init(self.conv, options.initializer)
def __init__(self, component, args, options):
super(EfficientAttention, self).__init__(component, args, options)
self.dims = list(component.current_size.dims).copy()
in_dim = self.dims[0]
out_dim = in_dim
if (len(args) > 0):
out_dim = args[0]
self.in_channels = in_dim
self.key_channels = options.key_channels or 16
self.head_count = options.heads or 4
self.value_channels = options.value_channels or 16
self.keys = nn.Conv2d(in_dim, self.key_channels, 1)
self.queries = nn.Conv2d(in_dim, self.key_channels, 1)
self.values = nn.Conv2d(in_dim, self.value_channels, 1)
self.reprojection = nn.Conv2d(self.value_channels, out_dim, 1)
self.size = LayerShape(out_dim, self.dims[1], self.dims[2])
def _sample(self):
gan = self.gan
gan.x = self.bw
self.xstep += 1
if self.xstep > self.xstep_count:
self.x = gan.inputs.next()
#self.x = torch.unsqueeze(self.gan.x[0],0).repeat(gan.batch_size(),1,1,1)
self.bw = BW(gan, None, None, hc.Config({}),
LayerShape(*self.x.shape[1:])).forward_grayscale(
self.x).repeat(1, 3, 1, 1)
self.xstep = 0
self.pos = self.direction * self.velocity + self.pos
self.gan.latent.z = self.pos
mask = torch.gt(self.pos, self.ones)
mask += torch.lt(self.pos, -self.ones)
self.direction = self.direction + 2 * self.direction * (-self.ones *
mask)
g = gan.generator.forward(self.pos)
return [('generator', g)]
def output_size(self):
return LayerShape(*self.dims)
def output_size(self):
return LayerShape(self.channels, self.dims[1], self.dims[2])
def __init__(self, component, args, options):
super(Cat, self).__init__("cat", component, args, options)
self.size = LayerShape(
sum([layer.channels for layer in self.layer_sizes]),
*self.size.dims[1:])
def output_size(self):
return LayerShape(self.dims[0] // 4, self.dims[1] * 2,
self.dims[2] * 2)
def create(self):
self.latent = self.create_component("latent")
self.x = self.inputs.next()[0]
if self.config.ali:
self.encoder = self.create_component("encoder", input=self.x, context_shapes={"y": LayerShape(1)})
self.generator = self.create_component("generator", input=self.encoder, context_shapes={"y": LayerShape(1)})
self.discriminator = self.create_component("discriminator", context_shapes={"z": self.encoder.layer_shape()})
else:
self.generator = self.create_component("generator", input=self.x, context_shapes={"y": LayerShape(1)})
self.discriminator = self.create_component("discriminator")
self.classification = 0