DeepFaceLab/nnlib/CAInitializer.py
2019-03-19 23:53:27 +04:00

113 lines
3.6 KiB
Python

import numpy as np
def _compute_fans(shape, data_format='channels_last'):
"""Computes the number of input and output units for a weight shape.
# Arguments
shape: Integer shape tuple.
data_format: Image data format to use for convolution kernels.
Note that all kernels in Keras are standardized on the
`channels_last` ordering (even when inputs are set
to `channels_first`).
# Returns
A tuple of scalars, `(fan_in, fan_out)`.
# Raises
ValueError: in case of invalid `data_format` argument.
"""
if len(shape) == 2:
fan_in = shape[0]
fan_out = shape[1]
elif len(shape) in {3, 4, 5}:
# Assuming convolution kernels (1D, 2D or 3D).
# TH kernel shape: (depth, input_depth, ...)
# TF kernel shape: (..., input_depth, depth)
if data_format == 'channels_first':
receptive_field_size = np.prod(shape[2:])
fan_in = shape[1] * receptive_field_size
fan_out = shape[0] * receptive_field_size
elif data_format == 'channels_last':
receptive_field_size = np.prod(shape[:-2])
fan_in = shape[-2] * receptive_field_size
fan_out = shape[-1] * receptive_field_size
else:
raise ValueError('Invalid data_format: ' + data_format)
else:
# No specific assumptions.
fan_in = np.sqrt(np.prod(shape))
fan_out = np.sqrt(np.prod(shape))
return fan_in, fan_out
def _create_basis(filters, size, floatx, eps_std):
if size == 1:
return np.random.normal(0.0, eps_std, (filters, size))
nbb = filters // size + 1
li = []
for i in range(nbb):
a = np.random.normal(0.0, 1.0, (size, size))
a = _symmetrize(a)
u, _, v = np.linalg.svd(a)
li.extend(u.T.tolist())
p = np.array(li[:filters], dtype=floatx)
return p
def _symmetrize(a):
return a + a.T - np.diag(a.diagonal())
def _scale_filters(filters, variance):
c_var = np.var(filters)
p = np.sqrt(variance / c_var)
return filters * p
def CAGenerateWeights ( shape, floatx, data_format, eps_std=0.05, seed=None ):
if seed is not None:
np.random.seed(seed)
fan_in, fan_out = _compute_fans(shape, data_format)
variance = 2 / fan_in
rank = len(shape)
if rank == 3:
row, stack_size, filters_size = shape
transpose_dimensions = (2, 1, 0)
kernel_shape = (row,)
correct_ifft = lambda shape, s=[None]: np.fft.irfft(shape, s[0])
correct_fft = np.fft.rfft
elif rank == 4:
row, column, stack_size, filters_size = shape
transpose_dimensions = (2, 3, 1, 0)
kernel_shape = (row, column)
correct_ifft = np.fft.irfft2
correct_fft = np.fft.rfft2
elif rank == 5:
x, y, z, stack_size, filters_size = shape
transpose_dimensions = (3, 4, 0, 1, 2)
kernel_shape = (x, y, z)
correct_fft = np.fft.rfftn
correct_ifft = np.fft.irfftn
else:
raise ValueError('rank unsupported')
kernel_fourier_shape = correct_fft(np.zeros(kernel_shape)).shape
init = []
for i in range(filters_size):
basis = _create_basis(
stack_size, np.prod(kernel_fourier_shape), floatx, eps_std)
basis = basis.reshape((stack_size,) + kernel_fourier_shape)
filters = [correct_ifft(x, kernel_shape) +
np.random.normal(0, eps_std, kernel_shape) for
x in basis]
init.append(filters)
# Format of array is now: filters, stack, row, column
init = np.array(init)
init = _scale_filters(init, variance)
return init.transpose(transpose_dimensions)