github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-08 05:51:40 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

changing SubpixelUpscaler to variable H,W dims,

Browse source 
tensorflow backend : using depth_to_space in SubpixelUpscaler, so training speed increased by 4%
This commit is contained in:
iperov 2019-03-28 17:55:42 +04:00
parent 689aefeb2f
commit 4683c362ac
1 changed files with 114 additions and 63 deletions
Download patch file Download diff file Expand all files Collapse all files

75
nnlib/nnlib.py
View file

@ -321,6 +321,63 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
nnlib.dssim = dssim nnlib.dssim = dssim
if 'tensorflow' in backend:
class PixelShuffler(keras.layers.Layer):
def __init__(self, size=(2, 2), data_format='channels_last', **kwargs):
super(PixelShuffler, self).__init__(**kwargs)
self.data_format = data_format
self.size = size
def call(self, inputs):
input_shape = K.shape(inputs)
if K.int_shape(input_shape)[0] != 4:
raise ValueError('Inputs should have rank 4; Received input shape:', str(K.int_shape(inputs)))
if self.data_format == 'channels_first':
return K.tf.depth_to_space(inputs, self.size[0], 'NCHW')
elif self.data_format == 'channels_last':
return K.tf.depth_to_space(inputs, self.size[0], 'NHWC')
def compute_output_shape(self, input_shape):
if len(input_shape) != 4:
raise ValueError('Inputs should have rank ' +
str(4) +
'; Received input shape:', str(input_shape))
if self.data_format == 'channels_first':
height = input_shape[2] * self.size[0] if input_shape[2] is not None else None
width = input_shape[3] * self.size[1] if input_shape[3] is not None else None
channels = input_shape[1] // self.size[0] // self.size[1]
if channels * self.size[0] * self.size[1] != input_shape[1]:
raise ValueError('channels of input and size are incompatible')
return (input_shape[0],
channels,
height,
width)
elif self.data_format == 'channels_last':
height = input_shape[1] * self.size[0] if input_shape[1] is not None else None
width = input_shape[2] * self.size[1] if input_shape[2] is not None else None
channels = input_shape[3] // self.size[0] // self.size[1]
if channels * self.size[0] * self.size[1] != input_shape[3]:
raise ValueError('channels of input and size are incompatible')
return (input_shape[0],
height,
width,
channels)
def get_config(self):
config = {'size': self.size,
'data_format': self.data_format}
base_config = super(PixelShuffler, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
else:
class PixelShuffler(KL.Layer): class PixelShuffler(KL.Layer):
def __init__(self, size=(2, 2), data_format='channels_last', **kwargs): def __init__(self, size=(2, 2), data_format='channels_last', **kwargs):
super(PixelShuffler, self).__init__(**kwargs) super(PixelShuffler, self).__init__(**kwargs)
@ -328,16 +385,13 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
self.size = size self.size = size
def call(self, inputs): def call(self, inputs):
input_shape = K.int_shape(inputs)
if len(input_shape) != 4: input_shape = K.shape(inputs)
raise ValueError('Inputs should have rank ' + if K.int_shape(input_shape)[0] != 4:
str(4) + raise ValueError('Inputs should have rank 4; Received input shape:', str(K.int_shape(inputs)))
'; Received input shape:', str(input_shape))
if self.data_format == 'channels_first': if self.data_format == 'channels_first':
batch_size, c, h, w = input_shape batch_size, c, h, w = input_shape[0], K.int_shape(inputs)[1], input_shape[2], input_shape[3]
if batch_size is None:
batch_size = -1
rh, rw = self.size rh, rw = self.size
oh, ow = h * rh, w * rw oh, ow = h * rh, w * rw
oc = c // (rh * rw) oc = c // (rh * rw)
@ -348,9 +402,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
return out return out
elif self.data_format == 'channels_last': elif self.data_format == 'channels_last':
batch_size, h, w, c = input_shape batch_size, h, w, c = input_shape[0], input_shape[1], input_shape[2], K.int_shape(inputs)[-1]
if batch_size is None:
batch_size = -1
rh, rw = self.size rh, rw = self.size
oh, ow = h * rh, w * rw oh, ow = h * rh, w * rw
oc = c // (rh * rw) oc = c // (rh * rw)
@ -361,7 +413,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
return out return out
def compute_output_shape(self, input_shape): def compute_output_shape(self, input_shape):
if len(input_shape) != 4: if len(input_shape) != 4:
raise ValueError('Inputs should have rank ' + raise ValueError('Inputs should have rank ' +
str(4) + str(4) +