mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 13:33:24 -07:00
global refactoring and fixes,
removed support of extracted(aligned) PNG faces. Use old builds to convert from PNG to JPG. fanseg model file in facelib/ is renamed
This commit is contained in:
Colombo
parent
921b464d5b
commit
61472cdaf7
82 changed files with 3838 additions and 3812 deletions
|
|
@ -17,4 +17,6 @@ from .common import normalize_channels, cut_odd_image, overlay_alpha_image
|
|||
|
||||
from .IEPolys import IEPolys
|
||||
|
||||
from .blur import LinearMotionBlur
|
||||
from .blursharpen import LinearMotionBlur, blursharpen
|
||||
|
||||
from .filters import apply_random_hsv_shift, apply_random_motion_blur, apply_random_gaussian_blur, apply_random_bilinear_resize
|
||||
|
|
|
|||
|
|
@ -1,9 +0,0 @@
|
|||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def LinearMotionBlur(image, size, angle):
|
||||
k = np.zeros((size, size), dtype=np.float32)
|
||||
k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32)
|
||||
k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) )
|
||||
k = k * ( 1.0 / np.sum(k) )
|
||||
return cv2.filter2D(image, -1, k)
|
||||
38
core/imagelib/blursharpen.py
Normal file
38
core/imagelib/blursharpen.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def LinearMotionBlur(image, size, angle):
|
||||
k = np.zeros((size, size), dtype=np.float32)
|
||||
k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32)
|
||||
k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) )
|
||||
k = k * ( 1.0 / np.sum(k) )
|
||||
return cv2.filter2D(image, -1, k)
|
||||
|
||||
def blursharpen (img, sharpen_mode=0, kernel_size=3, amount=100):
|
||||
if kernel_size % 2 == 0:
|
||||
kernel_size += 1
|
||||
if amount > 0:
|
||||
if sharpen_mode == 1: #box
|
||||
kernel = np.zeros( (kernel_size, kernel_size), dtype=np.float32)
|
||||
kernel[ kernel_size//2, kernel_size//2] = 1.0
|
||||
box_filter = np.ones( (kernel_size, kernel_size), dtype=np.float32) / (kernel_size**2)
|
||||
kernel = kernel + (kernel - box_filter) * amount
|
||||
return cv2.filter2D(img, -1, kernel)
|
||||
elif sharpen_mode == 2: #gaussian
|
||||
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
|
||||
img = cv2.addWeighted(img, 1.0 + (0.5 * amount), blur, -(0.5 * amount), 0)
|
||||
return img
|
||||
elif amount < 0:
|
||||
n = -amount
|
||||
while n > 0:
|
||||
|
||||
img_blur = cv2.medianBlur(img, 5)
|
||||
if int(n / 10) != 0:
|
||||
img = img_blur
|
||||
else:
|
||||
pass_power = (n % 10) / 10.0
|
||||
img = img*(1.0-pass_power)+img_blur*pass_power
|
||||
n = max(n-10,0)
|
||||
|
||||
return img
|
||||
return img
|
||||
53
core/imagelib/filters.py
Normal file
53
core/imagelib/filters.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
import numpy as np
|
||||
from .blursharpen import LinearMotionBlur
|
||||
import cv2
|
||||
|
||||
def apply_random_hsv_shift(img, rnd_state=None):
|
||||
if rnd_state is None:
|
||||
rnd_state = np.random
|
||||
|
||||
h, s, v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
|
||||
h = ( h + rnd_state.randint(360) ) % 360
|
||||
s = np.clip ( s + rnd_state.random()-0.5, 0, 1 )
|
||||
v = np.clip ( v + rnd_state.random()/2-0.25, 0, 1 )
|
||||
img = np.clip( cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2BGR) , 0, 1 )
|
||||
return img
|
||||
|
||||
def apply_random_motion_blur( img, chance, mb_max_size, rnd_state=None ):
|
||||
if rnd_state is None:
|
||||
rnd_state = np.random
|
||||
|
||||
mblur_rnd_kernel = rnd_state.randint(mb_max_size)+1
|
||||
mblur_rnd_deg = rnd_state.randint(360)
|
||||
|
||||
if rnd_state.randint(100) < np.clip(chance, 0, 100):
|
||||
img = LinearMotionBlur (img, mblur_rnd_kernel, mblur_rnd_deg )
|
||||
|
||||
return img
|
||||
|
||||
def apply_random_gaussian_blur( img, chance, kernel_max_size, rnd_state=None ):
|
||||
if rnd_state is None:
|
||||
rnd_state = np.random
|
||||
|
||||
if rnd_state.randint(100) < np.clip(chance, 0, 100):
|
||||
gblur_rnd_kernel = rnd_state.randint(kernel_max_size)*2+1
|
||||
img = cv2.GaussianBlur(img, (gblur_rnd_kernel,)*2 , 0)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def apply_random_bilinear_resize( img, chance, max_size_per, rnd_state=None ):
|
||||
if rnd_state is None:
|
||||
rnd_state = np.random
|
||||
|
||||
if rnd_state.randint(100) < np.clip(chance, 0, 100):
|
||||
h,w,c = img.shape
|
||||
|
||||
trg = rnd_state.rand()
|
||||
rw = w - int( trg * int(w*(max_size_per/100.0)) )
|
||||
rh = h - int( trg * int(h*(max_size_per/100.0)) )
|
||||
|
||||
img = cv2.resize (img, (rw,rh), cv2.INTER_LINEAR )
|
||||
img = cv2.resize (img, (w,h), cv2.INTER_LINEAR )
|
||||
|
||||
return img
|
||||
|
|
@ -1,13 +1,16 @@
|
|||
import multiprocessing
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
|
||||
import colorama
|
||||
import cv2
|
||||
from tqdm import tqdm
|
||||
|
||||
from core import stdex
|
||||
|
||||
try:
|
||||
import IPython #if success we are in colab
|
||||
from IPython.display import display, clear_output
|
||||
|
|
@ -38,6 +41,8 @@ class InteractBase(object):
|
|||
self.focus_wnd_name = None
|
||||
self.error_log_line_prefix = '/!\\ '
|
||||
|
||||
self.process_messages_callbacks = {}
|
||||
|
||||
def is_support_windows(self):
|
||||
return False
|
||||
|
||||
|
|
@ -164,7 +169,21 @@ class InteractBase(object):
|
|||
self.pg_bar.close()
|
||||
self.pg_bar = None
|
||||
|
||||
def add_process_messages_callback(self, func ):
|
||||
tid = threading.get_ident()
|
||||
callbacks = self.process_messages_callbacks.get(tid, None)
|
||||
if callbacks is None:
|
||||
callbacks = []
|
||||
self.process_messages_callbacks[tid] = callbacks
|
||||
|
||||
callbacks.append ( func )
|
||||
|
||||
def process_messages(self, sleep_time=0):
|
||||
callbacks = self.process_messages_callbacks.get(threading.get_ident(), None)
|
||||
if callbacks is not None:
|
||||
for func in callbacks:
|
||||
func()
|
||||
|
||||
self.on_process_messages(sleep_time)
|
||||
|
||||
def wait_any_key(self):
|
||||
|
|
@ -359,11 +378,11 @@ class InteractBase(object):
|
|||
def input_process(self, stdin_fd, sq, str):
|
||||
sys.stdin = os.fdopen(stdin_fd)
|
||||
try:
|
||||
inp = input (str)
|
||||
inp = input (str)
|
||||
sq.put (True)
|
||||
except:
|
||||
sq.put (False)
|
||||
|
||||
sq.put (False)
|
||||
|
||||
def input_in_time (self, str, max_time_sec):
|
||||
sq = multiprocessing.Queue()
|
||||
p = multiprocessing.Process(target=self.input_process, args=( sys.stdin.fileno(), sq, str))
|
||||
|
|
@ -377,14 +396,14 @@ class InteractBase(object):
|
|||
break
|
||||
if time.time() - t > max_time_sec:
|
||||
break
|
||||
|
||||
|
||||
p.terminate()
|
||||
|
||||
|
||||
p.terminate()
|
||||
p.join()
|
||||
|
||||
|
||||
old_stdin = sys.stdin
|
||||
sys.stdin = os.fdopen( os.dup(sys.stdin.fileno()) )
|
||||
old_stdin.close()
|
||||
old_stdin.close()
|
||||
return inp
|
||||
|
||||
def input_process_skip_pending(self, stdin_fd):
|
||||
|
|
|
|||
32
core/joblib/MPClassFuncOnDemand.py
Normal file
32
core/joblib/MPClassFuncOnDemand.py
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
import multiprocessing
|
||||
from core.interact import interact as io
|
||||
|
||||
class MPClassFuncOnDemand():
|
||||
def __init__(self, class_handle, class_func_name, **class_kwargs):
|
||||
self.class_handle = class_handle
|
||||
self.class_func_name = class_func_name
|
||||
self.class_kwargs = class_kwargs
|
||||
|
||||
self.class_func = None
|
||||
|
||||
self.s2c = multiprocessing.Queue()
|
||||
self.c2s = multiprocessing.Queue()
|
||||
self.lock = multiprocessing.Lock()
|
||||
|
||||
io.add_process_messages_callback(self.io_callback)
|
||||
|
||||
def io_callback(self):
|
||||
while not self.c2s.empty():
|
||||
func_args, func_kwargs = self.c2s.get()
|
||||
if self.class_func is None:
|
||||
self.class_func = getattr( self.class_handle(**self.class_kwargs), self.class_func_name)
|
||||
self.s2c.put ( self.class_func (*func_args, **func_kwargs) )
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
with self.lock:
|
||||
self.c2s.put ( (args, kwargs) )
|
||||
return self.s2c.get()
|
||||
|
||||
def __getstate__(self):
|
||||
return {'s2c':self.s2c, 'c2s':self.c2s, 'lock':self.lock}
|
||||
|
||||
25
core/joblib/MPFunc.py
Normal file
25
core/joblib/MPFunc.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
import multiprocessing
|
||||
from core.interact import interact as io
|
||||
|
||||
class MPFunc():
|
||||
def __init__(self, func):
|
||||
self.func = func
|
||||
|
||||
self.s2c = multiprocessing.Queue()
|
||||
self.c2s = multiprocessing.Queue()
|
||||
self.lock = multiprocessing.Lock()
|
||||
|
||||
io.add_process_messages_callback(self.io_callback)
|
||||
|
||||
def io_callback(self):
|
||||
while not self.c2s.empty():
|
||||
func_args, func_kwargs = self.c2s.get()
|
||||
self.s2c.put ( self.func (*func_args, **func_kwargs) )
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
with self.lock:
|
||||
self.c2s.put ( (args, kwargs) )
|
||||
return self.s2c.get()
|
||||
|
||||
def __getstate__(self):
|
||||
return {'s2c':self.s2c, 'c2s':self.c2s, 'lock':self.lock}
|
||||
|
|
@ -1,49 +0,0 @@
|
|||
import time
|
||||
import multiprocessing
|
||||
|
||||
class SubprocessFunctionCaller(object):
|
||||
class CliFunction(object):
|
||||
def __init__(self, s2c, c2s, lock):
|
||||
self.s2c = s2c
|
||||
self.c2s = c2s
|
||||
self.lock = lock
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
self.lock.acquire()
|
||||
self.c2s.put ( {'args':args, 'kwargs':kwargs} )
|
||||
while True:
|
||||
if not self.s2c.empty():
|
||||
obj = self.s2c.get()
|
||||
self.lock.release()
|
||||
return obj
|
||||
time.sleep(0.005)
|
||||
|
||||
class HostProcessor(object):
|
||||
def __init__(self, s2c, c2s, func):
|
||||
self.s2c = s2c
|
||||
self.c2s = c2s
|
||||
self.func = func
|
||||
|
||||
def process_messages(self):
|
||||
while not self.c2s.empty():
|
||||
obj = self.c2s.get()
|
||||
result = self.func ( *obj['args'], **obj['kwargs'] )
|
||||
self.s2c.put (result)
|
||||
|
||||
def __getstate__(self):
|
||||
#disable pickling this class
|
||||
return dict()
|
||||
|
||||
def __setstate__(self, d):
|
||||
self.__dict__.update(d)
|
||||
|
||||
@staticmethod
|
||||
def make_pair(func):
|
||||
s2c = multiprocessing.Queue()
|
||||
c2s = multiprocessing.Queue()
|
||||
lock = multiprocessing.Lock()
|
||||
|
||||
host_processor = SubprocessFunctionCaller.HostProcessor (s2c, c2s, func)
|
||||
cli_func = SubprocessFunctionCaller.CliFunction (s2c, c2s, lock)
|
||||
|
||||
return host_processor, cli_func
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
from .SubprocessorBase import Subprocessor
|
||||
from .SubprocessFunctionCaller import SubprocessFunctionCaller
|
||||
from .ThisThreadGenerator import ThisThreadGenerator
|
||||
from .SubprocessGenerator import SubprocessGenerator
|
||||
from .SubprocessGenerator import SubprocessGenerator
|
||||
from .MPFunc import MPFunc
|
||||
from .MPClassFuncOnDemand import MPClassFuncOnDemand
|
||||
|
|
@ -1,569 +0,0 @@
|
|||
|
||||
def initialize_archis(nn):
|
||||
tf = nn.tf
|
||||
|
||||
def get_ae_models(resolution):
|
||||
lowest_dense_res = resolution // 16
|
||||
conv_kernel_initializer = nn.initializers.ca()
|
||||
|
||||
class Downscale(nn.ModelBase):
|
||||
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.dilations = dilations
|
||||
self.subpixel = subpixel
|
||||
self.use_activator = use_activator
|
||||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv1 = nn.Conv2D( self.in_ch,
|
||||
self.out_ch // (4 if self.subpixel else 1),
|
||||
kernel_size=self.kernel_size,
|
||||
strides=1 if self.subpixel else 2,
|
||||
padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
if self.subpixel:
|
||||
x = nn.tf_space_to_depth(x, 2)
|
||||
if self.use_activator:
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return (self.out_ch // 4) * 4
|
||||
|
||||
class DownscaleBlock(nn.ModelBase):
|
||||
def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
|
||||
self.downs = []
|
||||
|
||||
last_ch = in_ch
|
||||
for i in range(n_downscales):
|
||||
cur_ch = ch*( min(2**i, 8) )
|
||||
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
|
||||
last_ch = self.downs[-1].get_out_ch()
|
||||
|
||||
def forward(self, inp):
|
||||
x = inp
|
||||
for down in self.downs:
|
||||
x = down(x)
|
||||
return x
|
||||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
x = nn.tf_depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
x = self.conv2(x)
|
||||
x = tf.nn.leaky_relu(inp + x, 0.2)
|
||||
return x
|
||||
|
||||
class UpdownResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, inner_ch, kernel_size=3 ):
|
||||
self.up = Upscale (ch, inner_ch, kernel_size=kernel_size)
|
||||
self.res = ResidualBlock (inner_ch, kernel_size=kernel_size)
|
||||
self.down = Downscale (inner_ch, ch, kernel_size=kernel_size, use_activator=False)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.up(inp)
|
||||
x = upx = self.res(x)
|
||||
x = self.down(x)
|
||||
x = x + inp
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
return x, upx
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, e_ch, is_hd):
|
||||
self.is_hd=is_hd
|
||||
if self.is_hd:
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=3, dilations=1)
|
||||
self.down2 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=5, dilations=1)
|
||||
self.down3 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=5, dilations=2)
|
||||
self.down4 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=7, dilations=2)
|
||||
else:
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5, dilations=1, subpixel=False)
|
||||
|
||||
def forward(self, inp):
|
||||
if self.is_hd:
|
||||
x = tf.concat([ nn.tf_flatten(self.down1(inp)),
|
||||
nn.tf_flatten(self.down2(inp)),
|
||||
nn.tf_flatten(self.down3(inp)),
|
||||
nn.tf_flatten(self.down4(inp)) ], -1 )
|
||||
else:
|
||||
x = nn.tf_flatten(self.down1(inp))
|
||||
return x
|
||||
|
||||
class Inter(nn.ModelBase):
|
||||
def __init__(self, in_ch, ae_ch, ae_out_ch, is_hd=False, **kwargs):
|
||||
self.in_ch, self.ae_ch, self.ae_out_ch = in_ch, ae_ch, ae_out_ch
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def on_build(self):
|
||||
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
|
||||
|
||||
self.dense1 = nn.Dense( in_ch, ae_ch )
|
||||
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
|
||||
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.dense1(inp)
|
||||
x = self.dense2(x)
|
||||
x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
|
||||
x = self.upscale1(x)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return self.ae_out_ch
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, d_ch, d_mask_ch, is_hd ):
|
||||
self.is_hd = is_hd
|
||||
|
||||
self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
|
||||
self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
|
||||
|
||||
if is_hd:
|
||||
self.res0 = UpdownResidualBlock(in_ch, d_ch*8, kernel_size=3)
|
||||
self.res1 = UpdownResidualBlock(d_ch*8, d_ch*4, kernel_size=3)
|
||||
self.res2 = UpdownResidualBlock(d_ch*4, d_ch*2, kernel_size=3)
|
||||
self.res3 = UpdownResidualBlock(d_ch*2, d_ch, kernel_size=3)
|
||||
else:
|
||||
self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
|
||||
self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
|
||||
self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
|
||||
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
|
||||
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
|
||||
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def get_weights_ex(self, include_mask):
|
||||
# Call internal get_weights in order to initialize inner logic
|
||||
self.get_weights()
|
||||
|
||||
weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() \
|
||||
+ self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.out_conv.get_weights()
|
||||
|
||||
if self.is_hd:
|
||||
weights += self.res3.get_weights()
|
||||
|
||||
if include_mask:
|
||||
weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() \
|
||||
+ self.out_convm.get_weights()
|
||||
return weights
|
||||
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
||||
if self.is_hd:
|
||||
x, upx = self.res0(z)
|
||||
x = self.upscale0(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res1(x)
|
||||
|
||||
x = self.upscale1(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res2(x)
|
||||
|
||||
x = self.upscale2(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res3(x)
|
||||
else:
|
||||
x = self.upscale0(z)
|
||||
x = self.res0(x)
|
||||
x = self.upscale1(x)
|
||||
x = self.res1(x)
|
||||
x = self.upscale2(x)
|
||||
x = self.res2(x)
|
||||
|
||||
m = self.upscalem0(z)
|
||||
m = self.upscalem1(m)
|
||||
m = self.upscalem2(m)
|
||||
|
||||
return tf.nn.sigmoid(self.out_conv(x)), \
|
||||
tf.nn.sigmoid(self.out_convm(m))
|
||||
|
||||
|
||||
|
||||
return lowest_dense_res, Encoder, Inter, Decoder
|
||||
|
||||
nn.get_ae_models = get_ae_models
|
||||
|
||||
def get_ae_models_chervonij(resolution):
|
||||
lowest_dense_res = resolution // 32
|
||||
"""
|
||||
by @chervonij
|
||||
"""
|
||||
conv_kernel_initializer = nn.initializers.ca()
|
||||
|
||||
class Downscale(nn.ModelBase):
|
||||
def __init__(self, in_ch, kernel_size=3, dilations=1, *kwargs ):
|
||||
self.in_ch = in_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.dilations = dilations
|
||||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv_base1 = nn.Conv2D( self.in_ch, self.in_ch//2, kernel_size=1, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
self.conv_l1 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
self.conv_l2 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
self.conv_base2 = nn.Conv2D( self.in_ch, self.in_ch//2, kernel_size=1, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
self.conv_r1 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
self.pool_size = [1,1,2,2] if nn.data_format == 'NCHW' else [1,2,2,1]
|
||||
def forward(self, x):
|
||||
|
||||
x_l = self.conv_base1(x)
|
||||
x_l = self.conv_l1(x_l)
|
||||
x_l = self.conv_l2(x_l)
|
||||
|
||||
x_r = self.conv_base2(x)
|
||||
x_r = self.conv_r1(x_r)
|
||||
|
||||
x_pool = tf.nn.max_pool(x, ksize=self.pool_size, strides=self.pool_size, padding='SAME', data_format=nn.data_format)
|
||||
|
||||
x = tf.concat([x_l, x_r, x_pool], axis=nn.conv2d_ch_axis)
|
||||
x = nn.tf_gelu(x)
|
||||
return x
|
||||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv2 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv3 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv4 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
x0 = self.conv1(x)
|
||||
x1 = self.conv2(x0)
|
||||
x2 = self.conv3(x1)
|
||||
x3 = self.conv4(x2)
|
||||
x = tf.concat([x0, x1, x2, x3], axis=nn.conv2d_ch_axis)
|
||||
x = nn.tf_gelu(x)
|
||||
x = nn.tf_depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
x = self.conv2(x)
|
||||
x = tf.nn.leaky_relu(inp + x, 0.2)
|
||||
return x
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, e_ch, **kwargs):
|
||||
self.conv0 = nn.Conv2D(in_ch, e_ch, kernel_size=3, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
self.down0 = Downscale(e_ch)
|
||||
self.down1 = Downscale(e_ch*2)
|
||||
self.down2 = Downscale(e_ch*4)
|
||||
self.down3 = Downscale(e_ch*8)
|
||||
self.down4 = Downscale(e_ch*16)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv0(inp)
|
||||
x = self.down0(x)
|
||||
x = self.down1(x)
|
||||
x = self.down2(x)
|
||||
x = self.down3(x)
|
||||
x = self.down4(x)
|
||||
x = nn.tf_flatten(x)
|
||||
return x
|
||||
|
||||
class Inter(nn.ModelBase):
|
||||
def __init__(self, in_ch, ae_ch, ae_out_ch, **kwargs):
|
||||
self.in_ch, self.ae_ch, self.ae_out_ch = in_ch, ae_ch, ae_out_ch
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def on_build(self, **kwargs):
|
||||
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
|
||||
|
||||
self.dense_l = nn.Dense( in_ch, ae_ch//2, kernel_initializer=tf.initializers.orthogonal)
|
||||
self.dense_r = nn.Dense( in_ch, ae_ch//2, maxout_features=4, kernel_initializer=tf.initializers.orthogonal)
|
||||
self.dense = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * (ae_out_ch//2), kernel_initializer=tf.initializers.orthogonal)
|
||||
self.upscale1 = Upscale(ae_out_ch//2, ae_out_ch//2)
|
||||
|
||||
def forward(self, inp):
|
||||
x0 = self.dense_l(inp)
|
||||
x1 = self.dense_r(inp)
|
||||
x = tf.concat([x0, x1], axis=-1)
|
||||
x = self.dense(x)
|
||||
x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch//2)
|
||||
x = self.upscale1(x)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return self.ae_out_ch//2
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, d_ch, d_mask_ch, **kwargs):
|
||||
|
||||
self.upscale0 = Upscale(in_ch, d_ch*8)
|
||||
self.upscale1 = Upscale(d_ch*8, d_ch*4)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2)
|
||||
self.upscale3 = Upscale(d_ch*2, d_ch)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
|
||||
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
|
||||
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
|
||||
self.upscalem3 = Upscale(d_mask_ch*2, d_mask_ch, kernel_size=3)
|
||||
self.out_convm = nn.Conv2D( d_mask_ch, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def get_weights_ex(self, include_mask):
|
||||
# Call internal get_weights in order to initialize inner logic
|
||||
self.get_weights()
|
||||
|
||||
weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() + self.upscale3.get_weights() + self.out_conv.get_weights()
|
||||
|
||||
if include_mask:
|
||||
weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() + self.upscale3.get_weights() + self.out_convm.get_weights()
|
||||
return weights
|
||||
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
||||
x = self.upscale0(inp)
|
||||
x = self.upscale1(x)
|
||||
x = self.upscale2(x)
|
||||
x = self.upscale3(x)
|
||||
|
||||
m = self.upscalem0(z)
|
||||
m = self.upscalem1(m)
|
||||
m = self.upscalem2(m)
|
||||
m = self.upscalem3(m)
|
||||
|
||||
return tf.nn.sigmoid(self.out_conv(x)), \
|
||||
tf.nn.sigmoid(self.out_convm(m))
|
||||
|
||||
return lowest_dense_res, Encoder, Inter, Decoder
|
||||
|
||||
nn.get_ae_models_chervonij = get_ae_models_chervonij
|
||||
|
||||
"""
|
||||
def get_ae_models2():
|
||||
conv_kernel_initializer = nn.initializers.ca()
|
||||
class Downscale(nn.ModelBase):
|
||||
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, is_hd=False, *kwargs ):
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.dilations = dilations
|
||||
self.is_hd = is_hd
|
||||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
if not self.is_hd:
|
||||
self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
else:
|
||||
self.conv0 = nn.Conv2D( self.in_ch, self.out_ch//4, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
self.conv1 = nn.Conv2D( self.out_ch//4, self.out_ch//4, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
self.conv2 = nn.Conv2D( self.out_ch//4, self.out_ch//4, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
self.conv3 = nn.Conv2D( self.out_ch//4, self.out_ch//4, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
if not self.is_hd:
|
||||
x = self.conv1(x)
|
||||
else:
|
||||
x = x0 = self.conv0(x)
|
||||
x = x1 = self.conv1(x)
|
||||
x = x2 = self.conv2(x)
|
||||
x = x3 = self.conv3(x)
|
||||
x = tf.concat([x0,x1,x2,x3], axis=nn.conv2d_ch_axis)
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return (self.out_ch // 4) * 4
|
||||
|
||||
class DownscaleBlock(nn.ModelBase):
|
||||
def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, is_hd=False):
|
||||
self.downs = []
|
||||
|
||||
last_ch = in_ch
|
||||
for i in range(n_downscales):
|
||||
cur_ch = ch*( min(2**i, 8) )
|
||||
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, is_hd=is_hd) )
|
||||
last_ch = self.downs[-1].get_out_ch()
|
||||
|
||||
def forward(self, inp):
|
||||
x = inp
|
||||
for down in self.downs:
|
||||
x = down(x)
|
||||
return x
|
||||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3, is_hd=False ):
|
||||
self.is_hd = is_hd
|
||||
|
||||
if not self.is_hd:
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
else:
|
||||
self.conv0 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv1 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv2 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv3 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
if not self.is_hd:
|
||||
x = self.conv1(x)
|
||||
else:
|
||||
x = x0 = self.conv0(x)
|
||||
x = x1 = self.conv1(x)
|
||||
x = x2 = self.conv2(x)
|
||||
x = x3 = self.conv3(x)
|
||||
x = tf.concat([x0,x1,x2,x3], axis=nn.conv2d_ch_axis)
|
||||
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
x = nn.tf_depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3, is_hd=False ):
|
||||
self.is_hd = is_hd
|
||||
if not is_hd:
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
else:
|
||||
self.conv10 = nn.Conv2D( ch, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv11 = nn.Conv2D( ch//4, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv12 = nn.Conv2D( ch//4, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.conv13 = nn.Conv2D( ch//4, ch//4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, inp):
|
||||
if not self.is_hd:
|
||||
x = self.conv1(inp)
|
||||
else:
|
||||
x = x0 = self.conv10(inp)
|
||||
x = x1 = self.conv11(x)
|
||||
x = x2 = self.conv12(x)
|
||||
x = x3 = self.conv13(x)
|
||||
x = tf.concat([x0,x1,x2,x3], axis=nn.conv2d_ch_axis)
|
||||
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
x = self.conv2(x)
|
||||
x = tf.nn.leaky_relu(inp + x, 0.2)
|
||||
return x
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, e_ch, is_hd):
|
||||
self.is_hd=is_hd
|
||||
if self.is_hd:
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=3, is_hd=is_hd)
|
||||
self.down2 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=5, is_hd=is_hd)
|
||||
else:
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5, is_hd=is_hd)
|
||||
|
||||
def forward(self, inp):
|
||||
if self.is_hd:
|
||||
x = tf.concat([ nn.tf_flatten(self.down1(inp)),
|
||||
nn.tf_flatten(self.down2(inp)) ], -1 )
|
||||
else:
|
||||
x = nn.tf_flatten(self.down1(inp))
|
||||
return x
|
||||
|
||||
class Inter(nn.ModelBase):
|
||||
def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, is_hd=False, **kwargs):
|
||||
self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch, self.is_hd = in_ch, lowest_dense_res, ae_ch, ae_out_ch, is_hd
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def on_build(self):
|
||||
in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
|
||||
|
||||
self.dense1 = nn.Dense( in_ch, ae_ch )
|
||||
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
|
||||
self.upscale1 = Upscale(ae_out_ch, ae_out_ch, is_hd=self.is_hd)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.dense1(inp)
|
||||
x = self.dense2(x)
|
||||
x = nn.tf_reshape_4D (x, self.lowest_dense_res, self.lowest_dense_res, self.ae_out_ch)
|
||||
x = self.upscale1(x)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return self.ae_out_ch
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, d_ch, d_mask_ch, is_hd ):
|
||||
self.is_hd = is_hd
|
||||
|
||||
self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
|
||||
self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
|
||||
|
||||
self.res0 = ResidualBlock(d_ch*8, kernel_size=3, is_hd=is_hd)
|
||||
self.res1 = ResidualBlock(d_ch*4, kernel_size=3, is_hd=is_hd)
|
||||
self.res2 = ResidualBlock(d_ch*2, kernel_size=3, is_hd=is_hd)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
|
||||
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
|
||||
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
|
||||
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def get_weights_ex(self, include_mask):
|
||||
# Call internal get_weights in order to initialize inner logic
|
||||
self.get_weights()
|
||||
|
||||
weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() \
|
||||
+ self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.out_conv.get_weights()
|
||||
|
||||
if include_mask:
|
||||
weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() \
|
||||
+ self.out_convm.get_weights()
|
||||
return weights
|
||||
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
||||
x = self.upscale0(z)
|
||||
x = self.res0(x)
|
||||
x = self.upscale1(x)
|
||||
x = self.res1(x)
|
||||
x = self.upscale2(x)
|
||||
x = self.res2(x)
|
||||
|
||||
m = self.upscalem0(z)
|
||||
m = self.upscalem1(m)
|
||||
m = self.upscalem2(m)
|
||||
|
||||
return tf.nn.sigmoid(self.out_conv(x)), \
|
||||
tf.nn.sigmoid(self.out_convm(m))
|
||||
|
||||
|
||||
return Encoder, Inter, Decoder
|
||||
|
||||
nn.get_ae_models2 = get_ae_models2
|
||||
"""
|
||||
17
core/leras/archis/ArchiBase.py
Normal file
17
core/leras/archis/ArchiBase.py
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
from core.leras import nn
|
||||
|
||||
class ArchiBase():
|
||||
|
||||
def __init__(self, *args, name=None, **kwargs):
|
||||
self.name=name
|
||||
|
||||
|
||||
#overridable
|
||||
def flow(self, *args, **kwargs):
|
||||
raise Exception("this archi does not support flow. Use model classes directly.")
|
||||
|
||||
#overridable
|
||||
def get_weights(self):
|
||||
pass
|
||||
|
||||
nn.ArchiBase = ArchiBase
|
||||
151
core/leras/archis/DFLSegnet.py
Normal file
151
core/leras/archis/DFLSegnet.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class DFLSegnetArchi(nn.ArchiBase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
class ConvBlock(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch):
|
||||
self.conv = nn.Conv2D (in_ch, out_ch, kernel_size=3, padding='SAME')
|
||||
self.frn = nn.FRNorm2D(out_ch)
|
||||
self.tlu = nn.TLU(out_ch)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.frn(x)
|
||||
x = self.tlu(x)
|
||||
return x
|
||||
|
||||
class UpConvBlock(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch):
|
||||
self.conv = nn.Conv2DTranspose (in_ch, out_ch, kernel_size=3, padding='SAME')
|
||||
self.frn = nn.FRNorm2D(out_ch)
|
||||
self.tlu = nn.TLU(out_ch)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.frn(x)
|
||||
x = self.tlu(x)
|
||||
return x
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, base_ch):
|
||||
self.conv01 = ConvBlock(in_ch, base_ch)
|
||||
self.conv02 = ConvBlock(base_ch, base_ch)
|
||||
self.bp0 = nn.BlurPool (filt_size=3)
|
||||
|
||||
|
||||
self.conv11 = ConvBlock(base_ch, base_ch*2)
|
||||
self.conv12 = ConvBlock(base_ch*2, base_ch*2)
|
||||
self.bp1 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.conv21 = ConvBlock(base_ch*2, base_ch*4)
|
||||
self.conv22 = ConvBlock(base_ch*4, base_ch*4)
|
||||
self.conv23 = ConvBlock(base_ch*4, base_ch*4)
|
||||
self.bp2 = nn.BlurPool (filt_size=3)
|
||||
|
||||
|
||||
self.conv31 = ConvBlock(base_ch*4, base_ch*8)
|
||||
self.conv32 = ConvBlock(base_ch*8, base_ch*8)
|
||||
self.conv33 = ConvBlock(base_ch*8, base_ch*8)
|
||||
self.bp3 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.conv41 = ConvBlock(base_ch*8, base_ch*8)
|
||||
self.conv42 = ConvBlock(base_ch*8, base_ch*8)
|
||||
self.conv43 = ConvBlock(base_ch*8, base_ch*8)
|
||||
self.bp4 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.conv_center = ConvBlock(base_ch*8, base_ch*8)
|
||||
|
||||
def forward(self, inp):
|
||||
x = inp
|
||||
|
||||
x = self.conv01(x)
|
||||
x = x0 = self.conv02(x)
|
||||
x = self.bp0(x)
|
||||
|
||||
x = self.conv11(x)
|
||||
x = x1 = self.conv12(x)
|
||||
x = self.bp1(x)
|
||||
|
||||
x = self.conv21(x)
|
||||
x = self.conv22(x)
|
||||
x = x2 = self.conv23(x)
|
||||
x = self.bp2(x)
|
||||
|
||||
x = self.conv31(x)
|
||||
x = self.conv32(x)
|
||||
x = x3 = self.conv33(x)
|
||||
x = self.bp3(x)
|
||||
|
||||
x = self.conv41(x)
|
||||
x = self.conv42(x)
|
||||
x = x4 = self.conv43(x)
|
||||
x = self.bp4(x)
|
||||
|
||||
x = self.conv_center(x)
|
||||
|
||||
return x0,x1,x2,x3,x4, x
|
||||
|
||||
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, base_ch, out_ch):
|
||||
|
||||
self.up4 = UpConvBlock (base_ch*8, base_ch*4)
|
||||
self.conv43 = ConvBlock(base_ch*12, base_ch*8)
|
||||
self.conv42 = ConvBlock(base_ch*8, base_ch*8)
|
||||
self.conv41 = ConvBlock(base_ch*8, base_ch*8)
|
||||
|
||||
self.up3 = UpConvBlock (base_ch*8, base_ch*4)
|
||||
self.conv33 = ConvBlock(base_ch*12, base_ch*8)
|
||||
self.conv32 = ConvBlock(base_ch*8, base_ch*8)
|
||||
self.conv31 = ConvBlock(base_ch*8, base_ch*8)
|
||||
|
||||
self.up2 = UpConvBlock (base_ch*8, base_ch*4)
|
||||
self.conv23 = ConvBlock(base_ch*8, base_ch*4)
|
||||
self.conv22 = ConvBlock(base_ch*4, base_ch*4)
|
||||
self.conv21 = ConvBlock(base_ch*4, base_ch*4)
|
||||
|
||||
self.up1 = UpConvBlock (base_ch*4, base_ch*2)
|
||||
self.conv12 = ConvBlock(base_ch*4, base_ch*2)
|
||||
self.conv11 = ConvBlock(base_ch*2, base_ch*2)
|
||||
|
||||
self.up0 = UpConvBlock (base_ch*2, base_ch)
|
||||
self.conv02 = ConvBlock(base_ch*2, base_ch)
|
||||
self.conv01 = ConvBlock(base_ch, base_ch)
|
||||
self.out_conv = nn.Conv2D (base_ch, out_ch, kernel_size=3, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x0,x1,x2,x3,x4,x = inp
|
||||
|
||||
x = self.up4(x)
|
||||
x = self.conv43(tf.concat([x,x4],axis=nn.conv2d_ch_axis))
|
||||
x = self.conv42(x)
|
||||
x = self.conv41(x)
|
||||
|
||||
x = self.up3(x)
|
||||
x = self.conv33(tf.concat([x,x3],axis=nn.conv2d_ch_axis))
|
||||
x = self.conv32(x)
|
||||
x = self.conv31(x)
|
||||
|
||||
x = self.up2(x)
|
||||
x = self.conv23(tf.concat([x,x2],axis=nn.conv2d_ch_axis))
|
||||
x = self.conv22(x)
|
||||
x = self.conv21(x)
|
||||
|
||||
x = self.up1(x)
|
||||
x = self.conv12(tf.concat([x,x1],axis=nn.conv2d_ch_axis))
|
||||
x = self.conv11(x)
|
||||
|
||||
x = self.up0(x)
|
||||
x = self.conv02(tf.concat([x,x0],axis=nn.conv2d_ch_axis))
|
||||
x = self.conv01(x)
|
||||
|
||||
logits = self.out_conv(x)
|
||||
return logits, tf.nn.sigmoid(logits)
|
||||
self.Encoder = Encoder
|
||||
self.Decoder = Decoder
|
||||
|
||||
nn.DFLSegnetArchi = DFLSegnetArchi
|
||||
490
core/leras/archis/DeepFakeArchi.py
Normal file
490
core/leras/archis/DeepFakeArchi.py
Normal file
|
|
@ -0,0 +1,490 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class DeepFakeArchi(nn.ArchiBase):
|
||||
"""
|
||||
resolution
|
||||
|
||||
mod None - default
|
||||
'chervonij'
|
||||
'quick'
|
||||
"""
|
||||
def __init__(self, resolution, mod=None):
|
||||
super().__init__()
|
||||
|
||||
if mod is None:
|
||||
class Downscale(nn.ModelBase):
|
||||
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.dilations = dilations
|
||||
self.subpixel = subpixel
|
||||
self.use_activator = use_activator
|
||||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv1 = nn.Conv2D( self.in_ch,
|
||||
self.out_ch // (4 if self.subpixel else 1),
|
||||
kernel_size=self.kernel_size,
|
||||
strides=1 if self.subpixel else 2,
|
||||
padding='SAME', dilations=self.dilations)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
if self.subpixel:
|
||||
x = nn.space_to_depth(x, 2)
|
||||
if self.use_activator:
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return (self.out_ch // 4) * 4
|
||||
|
||||
class DownscaleBlock(nn.ModelBase):
|
||||
def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
|
||||
self.downs = []
|
||||
|
||||
last_ch = in_ch
|
||||
for i in range(n_downscales):
|
||||
cur_ch = ch*( min(2**i, 8) )
|
||||
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
|
||||
last_ch = self.downs[-1].get_out_ch()
|
||||
|
||||
def forward(self, inp):
|
||||
x = inp
|
||||
for down in self.downs:
|
||||
x = down(x)
|
||||
return x
|
||||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
x = nn.depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
x = self.conv2(x)
|
||||
x = tf.nn.leaky_relu(inp + x, 0.2)
|
||||
return x
|
||||
|
||||
class UpdownResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, inner_ch, kernel_size=3 ):
|
||||
self.up = Upscale (ch, inner_ch, kernel_size=kernel_size)
|
||||
self.res = ResidualBlock (inner_ch, kernel_size=kernel_size)
|
||||
self.down = Downscale (inner_ch, ch, kernel_size=kernel_size, use_activator=False)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.up(inp)
|
||||
x = upx = self.res(x)
|
||||
x = self.down(x)
|
||||
x = x + inp
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
return x, upx
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, e_ch, is_hd):
|
||||
self.is_hd=is_hd
|
||||
if self.is_hd:
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=3, dilations=1)
|
||||
self.down2 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=5, dilations=1)
|
||||
self.down3 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=5, dilations=2)
|
||||
self.down4 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=7, dilations=2)
|
||||
else:
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5, dilations=1, subpixel=False)
|
||||
|
||||
def forward(self, inp):
|
||||
if self.is_hd:
|
||||
x = tf.concat([ nn.flatten(self.down1(inp)),
|
||||
nn.flatten(self.down2(inp)),
|
||||
nn.flatten(self.down3(inp)),
|
||||
nn.flatten(self.down4(inp)) ], -1 )
|
||||
else:
|
||||
x = nn.flatten(self.down1(inp))
|
||||
return x
|
||||
|
||||
lowest_dense_res = resolution // 16
|
||||
|
||||
class Inter(nn.ModelBase):
|
||||
def __init__(self, in_ch, ae_ch, ae_out_ch, is_hd=False, **kwargs):
|
||||
self.in_ch, self.ae_ch, self.ae_out_ch = in_ch, ae_ch, ae_out_ch
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def on_build(self):
|
||||
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
|
||||
|
||||
self.dense1 = nn.Dense( in_ch, ae_ch )
|
||||
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
|
||||
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.dense1(inp)
|
||||
x = self.dense2(x)
|
||||
x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
|
||||
x = self.upscale1(x)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def get_code_res():
|
||||
return lowest_dense_res
|
||||
|
||||
def get_out_ch(self):
|
||||
return self.ae_out_ch
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, d_ch, d_mask_ch, is_hd ):
|
||||
self.is_hd = is_hd
|
||||
|
||||
self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
|
||||
self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
|
||||
|
||||
if is_hd:
|
||||
self.res0 = UpdownResidualBlock(in_ch, d_ch*8, kernel_size=3)
|
||||
self.res1 = UpdownResidualBlock(d_ch*8, d_ch*4, kernel_size=3)
|
||||
self.res2 = UpdownResidualBlock(d_ch*4, d_ch*2, kernel_size=3)
|
||||
self.res3 = UpdownResidualBlock(d_ch*2, d_ch, kernel_size=3)
|
||||
else:
|
||||
self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
|
||||
self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
|
||||
self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME')
|
||||
|
||||
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
|
||||
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
|
||||
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
|
||||
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
||||
if self.is_hd:
|
||||
x, upx = self.res0(z)
|
||||
x = self.upscale0(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res1(x)
|
||||
|
||||
x = self.upscale1(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res2(x)
|
||||
|
||||
x = self.upscale2(x)
|
||||
x = tf.nn.leaky_relu(x + upx, 0.2)
|
||||
x, upx = self.res3(x)
|
||||
else:
|
||||
x = self.upscale0(z)
|
||||
x = self.res0(x)
|
||||
x = self.upscale1(x)
|
||||
x = self.res1(x)
|
||||
x = self.upscale2(x)
|
||||
x = self.res2(x)
|
||||
|
||||
m = self.upscalem0(z)
|
||||
m = self.upscalem1(m)
|
||||
m = self.upscalem2(m)
|
||||
|
||||
return tf.nn.sigmoid(self.out_conv(x)), \
|
||||
tf.nn.sigmoid(self.out_convm(m))
|
||||
|
||||
elif mod == 'chervonij':
|
||||
class Downscale(nn.ModelBase):
|
||||
def __init__(self, in_ch, kernel_size=3, dilations=1, *kwargs ):
|
||||
self.in_ch = in_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.dilations = dilations
|
||||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv_base1 = nn.Conv2D( self.in_ch, self.in_ch//2, kernel_size=1, strides=1, padding='SAME', dilations=self.dilations)
|
||||
self.conv_l1 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=1, padding='SAME', dilations=self.dilations)
|
||||
self.conv_l2 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations)
|
||||
|
||||
self.conv_base2 = nn.Conv2D( self.in_ch, self.in_ch//2, kernel_size=1, strides=1, padding='SAME', dilations=self.dilations)
|
||||
self.conv_r1 = nn.Conv2D( self.in_ch//2, self.in_ch//2, kernel_size=self.kernel_size, strides=2, padding='SAME', dilations=self.dilations)
|
||||
|
||||
self.pool_size = [1,1,2,2] if nn.data_format == 'NCHW' else [1,2,2,1]
|
||||
def forward(self, x):
|
||||
|
||||
x_l = self.conv_base1(x)
|
||||
x_l = self.conv_l1(x_l)
|
||||
x_l = self.conv_l2(x_l)
|
||||
|
||||
x_r = self.conv_base2(x)
|
||||
x_r = self.conv_r1(x_r)
|
||||
|
||||
x_pool = tf.nn.max_pool(x, ksize=self.pool_size, strides=self.pool_size, padding='SAME', data_format=nn.data_format)
|
||||
|
||||
x = tf.concat([x_l, x_r, x_pool], axis=nn.conv2d_ch_axis)
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
return x
|
||||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv3 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv4 = nn.Conv2D( out_ch, out_ch, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, x):
|
||||
x0 = self.conv1(x)
|
||||
x1 = self.conv2(x0)
|
||||
x2 = self.conv3(x1)
|
||||
x3 = self.conv4(x2)
|
||||
x = tf.concat([x0, x1, x2, x3], axis=nn.conv2d_ch_axis)
|
||||
x = tf.nn.leaky_relu(x, 0.1)
|
||||
x = nn.depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.norm = nn.FRNorm2D(ch)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
x = self.conv2(x)
|
||||
x = self.norm(inp + x)
|
||||
x = tf.nn.leaky_relu(x, 0.2)
|
||||
return x
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, e_ch, **kwargs):
|
||||
self.conv0 = nn.Conv2D(in_ch, e_ch, kernel_size=3, padding='SAME')
|
||||
|
||||
self.down0 = Downscale(e_ch)
|
||||
self.down1 = Downscale(e_ch*2)
|
||||
self.down2 = Downscale(e_ch*4)
|
||||
self.down3 = Downscale(e_ch*8)
|
||||
self.down4 = Downscale(e_ch*16)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv0(inp)
|
||||
x = self.down0(x)
|
||||
x = self.down1(x)
|
||||
x = self.down2(x)
|
||||
x = self.down3(x)
|
||||
x = self.down4(x)
|
||||
x = nn.flatten(x)
|
||||
return x
|
||||
|
||||
lowest_dense_res = resolution // 32
|
||||
|
||||
class Inter(nn.ModelBase):
|
||||
def __init__(self, in_ch, ae_ch, ae_out_ch, **kwargs):
|
||||
self.in_ch, self.ae_ch, self.ae_out_ch = in_ch, ae_ch, ae_out_ch
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def on_build(self, **kwargs):
|
||||
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
|
||||
|
||||
self.dense_l = nn.Dense( in_ch, ae_ch//2, kernel_initializer=tf.initializers.orthogonal)
|
||||
self.dense_r = nn.Dense( in_ch, ae_ch//2, kernel_initializer=tf.initializers.orthogonal)#maxout_ch=4,
|
||||
self.dense = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * (ae_out_ch//2), kernel_initializer=tf.initializers.orthogonal)
|
||||
self.upscale1 = Upscale(ae_out_ch//2, ae_out_ch//2)
|
||||
|
||||
def forward(self, inp):
|
||||
x0 = self.dense_l(inp)
|
||||
x1 = self.dense_r(inp)
|
||||
x = tf.concat([x0, x1], axis=-1)
|
||||
x = self.dense(x)
|
||||
x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch//2)
|
||||
x = self.upscale1(x)
|
||||
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return self.ae_out_ch//2
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, d_ch, d_mask_ch, **kwargs):
|
||||
|
||||
self.upscale0 = Upscale(in_ch, d_ch*8)
|
||||
self.upscale1 = Upscale(d_ch*8, d_ch*4)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2)
|
||||
self.upscale3 = Upscale(d_ch*2, d_ch)
|
||||
|
||||
self.res0 = ResidualBlock(d_ch*8)
|
||||
self.res1 = ResidualBlock(d_ch*4)
|
||||
self.res2 = ResidualBlock(d_ch*2)
|
||||
self.res3 = ResidualBlock(d_ch)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch, 3, kernel_size=1, padding='SAME')
|
||||
|
||||
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
|
||||
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
|
||||
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
|
||||
self.upscalem3 = Upscale(d_mask_ch*2, d_mask_ch, kernel_size=3)
|
||||
self.out_convm = nn.Conv2D( d_mask_ch, 1, kernel_size=1, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
||||
x = self.upscale0(z)
|
||||
x = self.res0(x)
|
||||
x = self.upscale1(x)
|
||||
x = self.res1(x)
|
||||
x = self.upscale2(x)
|
||||
x = self.res2(x)
|
||||
x = self.upscale3(x)
|
||||
x = self.res3(x)
|
||||
|
||||
m = self.upscalem0(z)
|
||||
m = self.upscalem1(m)
|
||||
m = self.upscalem2(m)
|
||||
m = self.upscalem3(m)
|
||||
|
||||
return tf.nn.sigmoid(self.out_conv(x)), \
|
||||
tf.nn.sigmoid(self.out_convm(m))
|
||||
elif mod == 'quick':
|
||||
class Downscale(nn.ModelBase):
|
||||
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.dilations = dilations
|
||||
self.subpixel = subpixel
|
||||
self.use_activator = use_activator
|
||||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv1 = nn.Conv2D( self.in_ch,
|
||||
self.out_ch // (4 if self.subpixel else 1),
|
||||
kernel_size=self.kernel_size,
|
||||
strides=1 if self.subpixel else 2,
|
||||
padding='SAME', dilations=self.dilations )
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
|
||||
if self.subpixel:
|
||||
x = nn.space_to_depth(x, 2)
|
||||
|
||||
if self.use_activator:
|
||||
x = nn.gelu(x)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return (self.out_ch // 4) * 4
|
||||
|
||||
class DownscaleBlock(nn.ModelBase):
|
||||
def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
|
||||
self.downs = []
|
||||
|
||||
last_ch = in_ch
|
||||
for i in range(n_downscales):
|
||||
cur_ch = ch*( min(2**i, 8) )
|
||||
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
|
||||
last_ch = self.downs[-1].get_out_ch()
|
||||
|
||||
def forward(self, inp):
|
||||
x = inp
|
||||
for down in self.downs:
|
||||
x = down(x)
|
||||
return x
|
||||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = nn.gelu(x)
|
||||
x = nn.depth_to_space(x, 2)
|
||||
return x
|
||||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
x = nn.gelu(x)
|
||||
x = self.conv2(x)
|
||||
x = inp + x
|
||||
x = nn.gelu(x)
|
||||
return x
|
||||
|
||||
class Encoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, e_ch):
|
||||
self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5)
|
||||
def forward(self, inp):
|
||||
return nn.flatten(self.down1(inp))
|
||||
|
||||
lowest_dense_res = resolution // 16
|
||||
|
||||
class Inter(nn.ModelBase):
|
||||
def __init__(self, in_ch, ae_ch, ae_out_ch, d_ch, **kwargs):
|
||||
self.in_ch, self.ae_ch, self.ae_out_ch, self.d_ch = in_ch, ae_ch, ae_out_ch, d_ch
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def on_build(self):
|
||||
in_ch, ae_ch, ae_out_ch, d_ch = self.in_ch, self.ae_ch, self.ae_out_ch, self.d_ch
|
||||
|
||||
self.dense1 = nn.Dense( in_ch, ae_ch, kernel_initializer=tf.initializers.orthogonal )
|
||||
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch, kernel_initializer=tf.initializers.orthogonal )
|
||||
self.upscale1 = Upscale(ae_out_ch, d_ch*8)
|
||||
self.res1 = ResidualBlock(d_ch*8)
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.dense1(inp)
|
||||
x = self.dense2(x)
|
||||
x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
|
||||
x = self.upscale1(x)
|
||||
x = self.res1(x)
|
||||
return x
|
||||
|
||||
def get_out_ch(self):
|
||||
return self.ae_out_ch
|
||||
|
||||
class Decoder(nn.ModelBase):
|
||||
def on_build(self, in_ch, d_ch):
|
||||
self.upscale1 = Upscale(in_ch, d_ch*4)
|
||||
self.res1 = ResidualBlock(d_ch*4)
|
||||
self.upscale2 = Upscale(d_ch*4, d_ch*2)
|
||||
self.res2 = ResidualBlock(d_ch*2)
|
||||
self.upscale3 = Upscale(d_ch*2, d_ch*1)
|
||||
self.res3 = ResidualBlock(d_ch*1)
|
||||
|
||||
self.upscalem1 = Upscale(in_ch, d_ch)
|
||||
self.upscalem2 = Upscale(d_ch, d_ch//2)
|
||||
self.upscalem3 = Upscale(d_ch//2, d_ch//2)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*1, 3, kernel_size=1, padding='SAME')
|
||||
self.out_convm = nn.Conv2D( d_ch//2, 1, kernel_size=1, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
x = self.upscale1 (z)
|
||||
x = self.res1 (x)
|
||||
x = self.upscale2 (x)
|
||||
x = self.res2 (x)
|
||||
x = self.upscale3 (x)
|
||||
x = self.res3 (x)
|
||||
|
||||
y = self.upscalem1 (z)
|
||||
y = self.upscalem2 (y)
|
||||
y = self.upscalem3 (y)
|
||||
|
||||
return tf.nn.sigmoid(self.out_conv(x)), \
|
||||
tf.nn.sigmoid(self.out_convm(y))
|
||||
|
||||
self.Encoder = Encoder
|
||||
self.Inter = Inter
|
||||
self.Decoder = Decoder
|
||||
|
||||
nn.DeepFakeArchi = DeepFakeArchi
|
||||
3
core/leras/archis/__init__.py
Normal file
3
core/leras/archis/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .ArchiBase import *
|
||||
from .DeepFakeArchi import *
|
||||
from .DFLSegnet import *
|
||||
|
|
@ -1,25 +1,6 @@
|
|||
import multiprocessing
|
||||
|
||||
import numpy as np
|
||||
|
||||
from core.joblib import Subprocessor
|
||||
|
||||
|
||||
def initialize_initializers(nn):
|
||||
tf = nn.tf
|
||||
from tensorflow.python.ops import init_ops
|
||||
|
||||
class initializers():
|
||||
class ca (init_ops.Initializer):
|
||||
def __call__(self, shape, dtype=None, partition_info=None):
|
||||
return tf.zeros( shape, dtype=dtype, name="_cai_")
|
||||
|
||||
@staticmethod
|
||||
def generate_batch( data_list, eps_std=0.05 ):
|
||||
# list of (shape, np.dtype)
|
||||
return CAInitializerSubprocessor (data_list).run()
|
||||
|
||||
nn.initializers = initializers
|
||||
import numpy as np
|
||||
|
||||
class CAInitializerSubprocessor(Subprocessor):
|
||||
@staticmethod
|
||||
|
|
@ -98,4 +79,4 @@ class CAInitializerSubprocessor(Subprocessor):
|
|||
|
||||
#override
|
||||
def get_result(self):
|
||||
return self.result
|
||||
return self.result
|
||||
20
core/leras/initializers/__init__.py
Normal file
20
core/leras/initializers/__init__.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
import numpy as np
|
||||
from tensorflow.python.ops import init_ops
|
||||
|
||||
from core.leras import nn
|
||||
|
||||
tf = nn.tf
|
||||
|
||||
from .CA import CAInitializerSubprocessor
|
||||
|
||||
class initializers():
|
||||
class ca (init_ops.Initializer):
|
||||
def __call__(self, shape, dtype=None, partition_info=None):
|
||||
return tf.zeros( shape, dtype=dtype, name="_cai_")
|
||||
|
||||
@staticmethod
|
||||
def generate_batch( data_list, eps_std=0.05 ):
|
||||
# list of (shape, np.dtype)
|
||||
return CAInitializerSubprocessor (data_list).run()
|
||||
|
||||
nn.initializers = initializers
|
||||
|
|
@ -1,573 +0,0 @@
|
|||
import pickle
|
||||
from pathlib import Path
|
||||
from core import pathex
|
||||
from core.interact import interact as io
|
||||
import numpy as np
|
||||
|
||||
|
||||
def initialize_layers(nn):
|
||||
tf = nn.tf
|
||||
|
||||
class Saveable():
|
||||
def __init__(self, name=None):
|
||||
self.name = name
|
||||
|
||||
#override
|
||||
def get_weights(self):
|
||||
#return tf tensors that should be initialized/loaded/saved
|
||||
pass
|
||||
|
||||
def save_weights(self, filename, force_dtype=None):
|
||||
d = {}
|
||||
weights = self.get_weights()
|
||||
|
||||
if self.name is None:
|
||||
raise Exception("name must be defined.")
|
||||
|
||||
name = self.name
|
||||
for w, w_val in zip(weights, nn.tf_sess.run (weights)):
|
||||
w_name_split = w.name.split('/', 1)
|
||||
if name != w_name_split[0]:
|
||||
raise Exception("weight first name != Saveable.name")
|
||||
|
||||
if force_dtype is not None:
|
||||
w_val = w_val.astype(force_dtype)
|
||||
|
||||
d[ w_name_split[1] ] = w_val
|
||||
|
||||
d_dumped = pickle.dumps (d, 4)
|
||||
pathex.write_bytes_safe ( Path(filename), d_dumped )
|
||||
|
||||
def load_weights(self, filename):
|
||||
"""
|
||||
returns True if file exists
|
||||
"""
|
||||
filepath = Path(filename)
|
||||
if filepath.exists():
|
||||
result = True
|
||||
d_dumped = filepath.read_bytes()
|
||||
d = pickle.loads(d_dumped)
|
||||
else:
|
||||
return False
|
||||
|
||||
weights = self.get_weights()
|
||||
|
||||
if self.name is None:
|
||||
raise Exception("name must be defined.")
|
||||
|
||||
tuples = []
|
||||
for w in weights:
|
||||
w_name_split = w.name.split('/')
|
||||
if self.name != w_name_split[0]:
|
||||
raise Exception("weight first name != Saveable.name")
|
||||
|
||||
sub_w_name = "/".join(w_name_split[1:])
|
||||
|
||||
w_val = d.get(sub_w_name, None)
|
||||
|
||||
if w_val is None:
|
||||
#io.log_err(f"Weight {w.name} was not loaded from file {filename}")
|
||||
tuples.append ( (w, w.initializer) )
|
||||
else:
|
||||
w_val = np.reshape( w_val, w.shape.as_list() )
|
||||
tuples.append ( (w, w_val) )
|
||||
|
||||
nn.tf_batch_set_value(tuples)
|
||||
|
||||
return True
|
||||
|
||||
def init_weights(self):
|
||||
nn.tf_init_weights(self.get_weights())
|
||||
|
||||
nn.Saveable = Saveable
|
||||
|
||||
class LayerBase():
|
||||
def __init__(self, name=None, **kwargs):
|
||||
self.name = name
|
||||
|
||||
#override
|
||||
def build_weights(self):
|
||||
pass
|
||||
|
||||
#override
|
||||
def get_weights(self):
|
||||
return []
|
||||
|
||||
def set_weights(self, new_weights):
|
||||
weights = self.get_weights()
|
||||
if len(weights) != len(new_weights):
|
||||
raise ValueError ('len of lists mismatch')
|
||||
|
||||
tuples = []
|
||||
for w, new_w in zip(weights, new_weights):
|
||||
if len(w.shape) != new_w.shape:
|
||||
new_w = new_w.reshape(w.shape)
|
||||
|
||||
tuples.append ( (w, new_w) )
|
||||
|
||||
nn.tf_batch_set_value (tuples)
|
||||
nn.LayerBase = LayerBase
|
||||
|
||||
class Conv2D(LayerBase):
|
||||
"""
|
||||
use_wscale bool enables equalized learning rate, if kernel_initializer is None, it will be forced to random_normal
|
||||
|
||||
|
||||
"""
|
||||
def __init__(self, in_ch, out_ch, kernel_size, strides=1, padding='SAME', dilations=1, use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
|
||||
if not isinstance(strides, int):
|
||||
raise ValueError ("strides must be an int type")
|
||||
if not isinstance(dilations, int):
|
||||
raise ValueError ("dilations must be an int type")
|
||||
kernel_size = int(kernel_size)
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.tf_floatx
|
||||
|
||||
if isinstance(padding, str):
|
||||
if padding == "SAME":
|
||||
padding = ( (kernel_size - 1) * dilations + 1 ) // 2
|
||||
elif padding == "VALID":
|
||||
padding = 0
|
||||
else:
|
||||
raise ValueError ("Wrong padding type. Should be VALID SAME or INT or 4x INTs")
|
||||
|
||||
if isinstance(padding, int):
|
||||
if padding != 0:
|
||||
if nn.data_format == "NHWC":
|
||||
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
|
||||
else:
|
||||
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
|
||||
else:
|
||||
padding = None
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
strides = [1,strides,strides,1]
|
||||
else:
|
||||
strides = [1,1,strides,strides]
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
dilations = [1,dilations,dilations,1]
|
||||
else:
|
||||
dilations = [1,1,dilations,dilations]
|
||||
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.strides = strides
|
||||
self.padding = padding
|
||||
self.dilations = dilations
|
||||
self.use_bias = use_bias
|
||||
self.use_wscale = use_wscale
|
||||
self.kernel_initializer = kernel_initializer
|
||||
self.bias_initializer = bias_initializer
|
||||
self.trainable = trainable
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = self.kernel_initializer
|
||||
if self.use_wscale:
|
||||
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
|
||||
fan_in = self.kernel_size*self.kernel_size*self.in_ch
|
||||
he_std = gain / np.sqrt(fan_in) # He init
|
||||
self.wscale = tf.constant(he_std, dtype=self.dtype )
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
|
||||
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
|
||||
|
||||
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.out_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
|
||||
|
||||
if self.use_bias:
|
||||
bias_initializer = self.bias_initializer
|
||||
if bias_initializer is None:
|
||||
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
|
||||
|
||||
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
|
||||
|
||||
def get_weights(self):
|
||||
weights = [self.weight]
|
||||
if self.use_bias:
|
||||
weights += [self.bias]
|
||||
return weights
|
||||
|
||||
def __call__(self, x):
|
||||
weight = self.weight
|
||||
if self.use_wscale:
|
||||
weight = weight * self.wscale
|
||||
|
||||
if self.padding is not None:
|
||||
x = tf.pad (x, self.padding, mode='CONSTANT')
|
||||
|
||||
x = tf.nn.conv2d(x, weight, self.strides, 'VALID', dilations=self.dilations, data_format=nn.data_format)
|
||||
if self.use_bias:
|
||||
if nn.data_format == "NHWC":
|
||||
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
|
||||
else:
|
||||
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
|
||||
x = tf.add(x, bias)
|
||||
return x
|
||||
|
||||
def __str__(self):
|
||||
r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
|
||||
|
||||
return r
|
||||
nn.Conv2D = Conv2D
|
||||
|
||||
class Conv2DTranspose(LayerBase):
|
||||
"""
|
||||
use_wscale enables weight scale (equalized learning rate)
|
||||
if kernel_initializer is None, it will be forced to random_normal
|
||||
"""
|
||||
def __init__(self, in_ch, out_ch, kernel_size, strides=2, padding='SAME', use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
|
||||
if not isinstance(strides, int):
|
||||
raise ValueError ("strides must be an int type")
|
||||
kernel_size = int(kernel_size)
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.tf_floatx
|
||||
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.strides = strides
|
||||
self.padding = padding
|
||||
self.use_bias = use_bias
|
||||
self.use_wscale = use_wscale
|
||||
self.kernel_initializer = kernel_initializer
|
||||
self.bias_initializer = bias_initializer
|
||||
self.trainable = trainable
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = self.kernel_initializer
|
||||
if self.use_wscale:
|
||||
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
|
||||
fan_in = self.kernel_size*self.kernel_size*self.in_ch
|
||||
he_std = gain / np.sqrt(fan_in) # He init
|
||||
self.wscale = tf.constant(he_std, dtype=self.dtype )
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
|
||||
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
|
||||
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.out_ch,self.in_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
|
||||
|
||||
if self.use_bias:
|
||||
bias_initializer = self.bias_initializer
|
||||
if bias_initializer is None:
|
||||
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
|
||||
|
||||
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
|
||||
|
||||
def get_weights(self):
|
||||
weights = [self.weight]
|
||||
if self.use_bias:
|
||||
weights += [self.bias]
|
||||
return weights
|
||||
|
||||
def __call__(self, x):
|
||||
shape = x.shape
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
h,w,c = shape[1], shape[2], shape[3]
|
||||
output_shape = tf.stack ( (tf.shape(x)[0],
|
||||
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
|
||||
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
|
||||
self.out_ch) )
|
||||
|
||||
strides = [1,self.strides,self.strides,1]
|
||||
else:
|
||||
c,h,w = shape[1], shape[2], shape[3]
|
||||
output_shape = tf.stack ( (tf.shape(x)[0],
|
||||
self.out_ch,
|
||||
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
|
||||
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
|
||||
) )
|
||||
strides = [1,1,self.strides,self.strides]
|
||||
weight = self.weight
|
||||
if self.use_wscale:
|
||||
weight = weight * self.wscale
|
||||
|
||||
x = tf.nn.conv2d_transpose(x, weight, output_shape, strides, padding=self.padding, data_format=nn.data_format)
|
||||
|
||||
if self.use_bias:
|
||||
if nn.data_format == "NHWC":
|
||||
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
|
||||
else:
|
||||
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
|
||||
x = tf.add(x, bias)
|
||||
return x
|
||||
|
||||
def __str__(self):
|
||||
r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
|
||||
|
||||
return r
|
||||
|
||||
def deconv_length(self, dim_size, stride_size, kernel_size, padding):
|
||||
assert padding in {'SAME', 'VALID', 'FULL'}
|
||||
if dim_size is None:
|
||||
return None
|
||||
if padding == 'VALID':
|
||||
dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0)
|
||||
elif padding == 'FULL':
|
||||
dim_size = dim_size * stride_size - (stride_size + kernel_size - 2)
|
||||
elif padding == 'SAME':
|
||||
dim_size = dim_size * stride_size
|
||||
return dim_size
|
||||
nn.Conv2DTranspose = Conv2DTranspose
|
||||
|
||||
class BlurPool(LayerBase):
|
||||
def __init__(self, filt_size=3, stride=2, **kwargs ):
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
self.strides = [1,stride,stride,1]
|
||||
else:
|
||||
self.strides = [1,1,stride,stride]
|
||||
|
||||
self.filt_size = filt_size
|
||||
pad = [ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ]
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
self.padding = [ [0,0], pad, pad, [0,0] ]
|
||||
else:
|
||||
self.padding = [ [0,0], [0,0], pad, pad ]
|
||||
|
||||
if(self.filt_size==1):
|
||||
a = np.array([1.,])
|
||||
elif(self.filt_size==2):
|
||||
a = np.array([1., 1.])
|
||||
elif(self.filt_size==3):
|
||||
a = np.array([1., 2., 1.])
|
||||
elif(self.filt_size==4):
|
||||
a = np.array([1., 3., 3., 1.])
|
||||
elif(self.filt_size==5):
|
||||
a = np.array([1., 4., 6., 4., 1.])
|
||||
elif(self.filt_size==6):
|
||||
a = np.array([1., 5., 10., 10., 5., 1.])
|
||||
elif(self.filt_size==7):
|
||||
a = np.array([1., 6., 15., 20., 15., 6., 1.])
|
||||
|
||||
a = a[:,None]*a[None,:]
|
||||
a = a / np.sum(a)
|
||||
a = a[:,:,None,None]
|
||||
self.a = a
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
self.k = tf.constant (self.a, dtype=nn.tf_floatx )
|
||||
|
||||
def __call__(self, x):
|
||||
k = tf.tile (self.k, (1,1,x.shape[nn.conv2d_ch_axis],1) )
|
||||
x = tf.pad(x, self.padding )
|
||||
x = tf.nn.depthwise_conv2d(x, k, self.strides, 'VALID', data_format=nn.data_format)
|
||||
return x
|
||||
nn.BlurPool = BlurPool
|
||||
|
||||
class Dense(LayerBase):
|
||||
def __init__(self, in_ch, out_ch, use_bias=True, use_wscale=False, maxout_ch=0, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
|
||||
"""
|
||||
use_wscale enables weight scale (equalized learning rate)
|
||||
if kernel_initializer is None, it will be forced to random_normal
|
||||
|
||||
maxout_ch https://link.springer.com/article/10.1186/s40537-019-0233-0
|
||||
typical 2-4 if you want to enable DenseMaxout behaviour
|
||||
"""
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.use_bias = use_bias
|
||||
self.use_wscale = use_wscale
|
||||
self.maxout_ch = maxout_ch
|
||||
self.kernel_initializer = kernel_initializer
|
||||
self.bias_initializer = bias_initializer
|
||||
self.trainable = trainable
|
||||
if dtype is None:
|
||||
dtype = nn.tf_floatx
|
||||
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
if self.maxout_ch > 1:
|
||||
weight_shape = (self.in_ch,self.out_ch*self.maxout_ch)
|
||||
else:
|
||||
weight_shape = (self.in_ch,self.out_ch)
|
||||
|
||||
kernel_initializer = self.kernel_initializer
|
||||
|
||||
if self.use_wscale:
|
||||
gain = 1.0
|
||||
fan_in = np.prod( weight_shape[:-1] )
|
||||
he_std = gain / np.sqrt(fan_in) # He init
|
||||
self.wscale = tf.constant(he_std, dtype=self.dtype )
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
|
||||
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
|
||||
|
||||
self.weight = tf.get_variable("weight", weight_shape, dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
|
||||
|
||||
if self.use_bias:
|
||||
bias_initializer = self.bias_initializer
|
||||
if bias_initializer is None:
|
||||
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
|
||||
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
|
||||
|
||||
def get_weights(self):
|
||||
weights = [self.weight]
|
||||
if self.use_bias:
|
||||
weights += [self.bias]
|
||||
return weights
|
||||
|
||||
def __call__(self, x):
|
||||
weight = self.weight
|
||||
if self.use_wscale:
|
||||
weight = weight * self.wscale
|
||||
|
||||
x = tf.matmul(x, weight)
|
||||
|
||||
if self.maxout_ch > 1:
|
||||
x = tf.reshape (x, (-1, self.out_ch, self.maxout_ch) )
|
||||
x = tf.reduce_max(x, axis=-1)
|
||||
|
||||
if self.use_bias:
|
||||
x = tf.add(x, tf.reshape(self.bias, (1,self.out_ch) ) )
|
||||
|
||||
return x
|
||||
nn.Dense = Dense
|
||||
|
||||
class InstanceNorm2D(LayerBase):
|
||||
def __init__(self, in_ch, dtype=None, **kwargs):
|
||||
self.in_ch = in_ch
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.tf_floatx
|
||||
self.dtype = dtype
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
|
||||
self.weight = tf.get_variable("weight", (self.in_ch,), dtype=self.dtype, initializer=kernel_initializer )
|
||||
self.bias = tf.get_variable("bias", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight, self.bias]
|
||||
|
||||
def __call__(self, x):
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (1,1,1,self.in_ch)
|
||||
else:
|
||||
shape = (1,self.in_ch,1,1)
|
||||
|
||||
weight = tf.reshape ( self.weight , shape )
|
||||
bias = tf.reshape ( self.bias , shape )
|
||||
|
||||
x_mean = tf.reduce_mean(x, axis=nn.conv2d_spatial_axes, keepdims=True )
|
||||
x_std = tf.math.reduce_std(x, axis=nn.conv2d_spatial_axes, keepdims=True ) + 1e-5
|
||||
|
||||
x = (x - x_mean) / x_std
|
||||
x *= weight
|
||||
x += bias
|
||||
|
||||
return x
|
||||
|
||||
nn.InstanceNorm2D = InstanceNorm2D
|
||||
|
||||
class BatchNorm2D(LayerBase):
|
||||
"""
|
||||
currently not for training
|
||||
"""
|
||||
def __init__(self, dim, eps=1e-05, momentum=0.1, dtype=None, **kwargs):
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
self.momentum = momentum
|
||||
if dtype is None:
|
||||
dtype = nn.tf_floatx
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
self.weight = tf.get_variable("weight", (self.dim,), dtype=self.dtype, initializer=tf.initializers.ones() )
|
||||
self.bias = tf.get_variable("bias", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros() )
|
||||
self.running_mean = tf.get_variable("running_mean", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
|
||||
self.running_var = tf.get_variable("running_var", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight, self.bias, self.running_mean, self.running_var]
|
||||
|
||||
def __call__(self, x):
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (1,1,1,self.dim)
|
||||
else:
|
||||
shape = (1,self.dim,1,1)
|
||||
|
||||
weight = tf.reshape ( self.weight , shape )
|
||||
bias = tf.reshape ( self.bias , shape )
|
||||
running_mean = tf.reshape ( self.running_mean, shape )
|
||||
running_var = tf.reshape ( self.running_var , shape )
|
||||
|
||||
x = (x - running_mean) / tf.sqrt( running_var + self.eps )
|
||||
x *= weight
|
||||
x += bias
|
||||
return x
|
||||
|
||||
nn.BatchNorm2D = BatchNorm2D
|
||||
|
||||
class AdaIN(LayerBase):
|
||||
"""
|
||||
"""
|
||||
def __init__(self, in_ch, mlp_ch, kernel_initializer=None, dtype=None, **kwargs):
|
||||
self.in_ch = in_ch
|
||||
self.mlp_ch = mlp_ch
|
||||
self.kernel_initializer = kernel_initializer
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.tf_floatx
|
||||
self.dtype = dtype
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = self.kernel_initializer
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.he_normal()#(dtype=self.dtype)
|
||||
|
||||
self.weight1 = tf.get_variable("weight1", (self.mlp_ch, self.in_ch), dtype=self.dtype, initializer=kernel_initializer)
|
||||
self.bias1 = tf.get_variable("bias1", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros())
|
||||
self.weight2 = tf.get_variable("weight2", (self.mlp_ch, self.in_ch), dtype=self.dtype, initializer=kernel_initializer)
|
||||
self.bias2 = tf.get_variable("bias2", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros())
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight1, self.bias1, self.weight2, self.bias2]
|
||||
|
||||
def __call__(self, inputs):
|
||||
x, mlp = inputs
|
||||
|
||||
gamma = tf.matmul(mlp, self.weight1)
|
||||
gamma = tf.add(gamma, tf.reshape(self.bias1, (1,self.in_ch) ) )
|
||||
|
||||
beta = tf.matmul(mlp, self.weight2)
|
||||
beta = tf.add(beta, tf.reshape(self.bias2, (1,self.in_ch) ) )
|
||||
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (-1,1,1,self.in_ch)
|
||||
else:
|
||||
shape = (-1,self.in_ch,1,1)
|
||||
|
||||
x_mean = tf.reduce_mean(x, axis=nn.conv2d_spatial_axes, keepdims=True )
|
||||
x_std = tf.math.reduce_std(x, axis=nn.conv2d_spatial_axes, keepdims=True ) + 1e-5
|
||||
|
||||
x = (x - x_mean) / x_std
|
||||
x *= tf.reshape(gamma, shape)
|
||||
|
||||
x += tf.reshape(beta, shape)
|
||||
|
||||
return x
|
||||
|
||||
nn.AdaIN = AdaIN
|
||||
56
core/leras/layers/AdaIN.py
Normal file
56
core/leras/layers/AdaIN.py
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class AdaIN(nn.LayerBase):
|
||||
"""
|
||||
"""
|
||||
def __init__(self, in_ch, mlp_ch, kernel_initializer=None, dtype=None, **kwargs):
|
||||
self.in_ch = in_ch
|
||||
self.mlp_ch = mlp_ch
|
||||
self.kernel_initializer = kernel_initializer
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
self.dtype = dtype
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = self.kernel_initializer
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.he_normal()
|
||||
|
||||
self.weight1 = tf.get_variable("weight1", (self.mlp_ch, self.in_ch), dtype=self.dtype, initializer=kernel_initializer)
|
||||
self.bias1 = tf.get_variable("bias1", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros())
|
||||
self.weight2 = tf.get_variable("weight2", (self.mlp_ch, self.in_ch), dtype=self.dtype, initializer=kernel_initializer)
|
||||
self.bias2 = tf.get_variable("bias2", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros())
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight1, self.bias1, self.weight2, self.bias2]
|
||||
|
||||
def forward(self, inputs):
|
||||
x, mlp = inputs
|
||||
|
||||
gamma = tf.matmul(mlp, self.weight1)
|
||||
gamma = tf.add(gamma, tf.reshape(self.bias1, (1,self.in_ch) ) )
|
||||
|
||||
beta = tf.matmul(mlp, self.weight2)
|
||||
beta = tf.add(beta, tf.reshape(self.bias2, (1,self.in_ch) ) )
|
||||
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (-1,1,1,self.in_ch)
|
||||
else:
|
||||
shape = (-1,self.in_ch,1,1)
|
||||
|
||||
x_mean = tf.reduce_mean(x, axis=nn.conv2d_spatial_axes, keepdims=True )
|
||||
x_std = tf.math.reduce_std(x, axis=nn.conv2d_spatial_axes, keepdims=True ) + 1e-5
|
||||
|
||||
x = (x - x_mean) / x_std
|
||||
x *= tf.reshape(gamma, shape)
|
||||
|
||||
x += tf.reshape(beta, shape)
|
||||
|
||||
return x
|
||||
|
||||
nn.AdaIN = AdaIN
|
||||
42
core/leras/layers/BatchNorm2D.py
Normal file
42
core/leras/layers/BatchNorm2D.py
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class BatchNorm2D(nn.LayerBase):
|
||||
"""
|
||||
currently not for training
|
||||
"""
|
||||
def __init__(self, dim, eps=1e-05, momentum=0.1, dtype=None, **kwargs):
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
self.momentum = momentum
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
self.weight = tf.get_variable("weight", (self.dim,), dtype=self.dtype, initializer=tf.initializers.ones() )
|
||||
self.bias = tf.get_variable("bias", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros() )
|
||||
self.running_mean = tf.get_variable("running_mean", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
|
||||
self.running_var = tf.get_variable("running_var", (self.dim,), dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight, self.bias, self.running_mean, self.running_var]
|
||||
|
||||
def forward(self, x):
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (1,1,1,self.dim)
|
||||
else:
|
||||
shape = (1,self.dim,1,1)
|
||||
|
||||
weight = tf.reshape ( self.weight , shape )
|
||||
bias = tf.reshape ( self.bias , shape )
|
||||
running_mean = tf.reshape ( self.running_mean, shape )
|
||||
running_var = tf.reshape ( self.running_var , shape )
|
||||
|
||||
x = (x - running_mean) / tf.sqrt( running_var + self.eps )
|
||||
x *= weight
|
||||
x += bias
|
||||
return x
|
||||
|
||||
nn.BatchNorm2D = BatchNorm2D
|
||||
50
core/leras/layers/BlurPool.py
Normal file
50
core/leras/layers/BlurPool.py
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
import numpy as np
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class BlurPool(nn.LayerBase):
|
||||
def __init__(self, filt_size=3, stride=2, **kwargs ):
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
self.strides = [1,stride,stride,1]
|
||||
else:
|
||||
self.strides = [1,1,stride,stride]
|
||||
|
||||
self.filt_size = filt_size
|
||||
pad = [ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ]
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
self.padding = [ [0,0], pad, pad, [0,0] ]
|
||||
else:
|
||||
self.padding = [ [0,0], [0,0], pad, pad ]
|
||||
|
||||
if(self.filt_size==1):
|
||||
a = np.array([1.,])
|
||||
elif(self.filt_size==2):
|
||||
a = np.array([1., 1.])
|
||||
elif(self.filt_size==3):
|
||||
a = np.array([1., 2., 1.])
|
||||
elif(self.filt_size==4):
|
||||
a = np.array([1., 3., 3., 1.])
|
||||
elif(self.filt_size==5):
|
||||
a = np.array([1., 4., 6., 4., 1.])
|
||||
elif(self.filt_size==6):
|
||||
a = np.array([1., 5., 10., 10., 5., 1.])
|
||||
elif(self.filt_size==7):
|
||||
a = np.array([1., 6., 15., 20., 15., 6., 1.])
|
||||
|
||||
a = a[:,None]*a[None,:]
|
||||
a = a / np.sum(a)
|
||||
a = a[:,:,None,None]
|
||||
self.a = a
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
self.k = tf.constant (self.a, dtype=nn.floatx )
|
||||
|
||||
def forward(self, x):
|
||||
k = tf.tile (self.k, (1,1,x.shape[nn.conv2d_ch_axis],1) )
|
||||
x = tf.pad(x, self.padding )
|
||||
x = tf.nn.depthwise_conv2d(x, k, self.strides, 'VALID', data_format=nn.data_format)
|
||||
return x
|
||||
nn.BlurPool = BlurPool
|
||||
112
core/leras/layers/Conv2D.py
Normal file
112
core/leras/layers/Conv2D.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
import numpy as np
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class Conv2D(nn.LayerBase):
|
||||
"""
|
||||
default kernel_initializer - CA
|
||||
use_wscale bool enables equalized learning rate, if kernel_initializer is None, it will be forced to random_normal
|
||||
|
||||
|
||||
"""
|
||||
def __init__(self, in_ch, out_ch, kernel_size, strides=1, padding='SAME', dilations=1, use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
|
||||
if not isinstance(strides, int):
|
||||
raise ValueError ("strides must be an int type")
|
||||
if not isinstance(dilations, int):
|
||||
raise ValueError ("dilations must be an int type")
|
||||
kernel_size = int(kernel_size)
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
|
||||
if isinstance(padding, str):
|
||||
if padding == "SAME":
|
||||
padding = ( (kernel_size - 1) * dilations + 1 ) // 2
|
||||
elif padding == "VALID":
|
||||
padding = 0
|
||||
else:
|
||||
raise ValueError ("Wrong padding type. Should be VALID SAME or INT or 4x INTs")
|
||||
|
||||
if isinstance(padding, int):
|
||||
if padding != 0:
|
||||
if nn.data_format == "NHWC":
|
||||
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
|
||||
else:
|
||||
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
|
||||
else:
|
||||
padding = None
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
strides = [1,strides,strides,1]
|
||||
else:
|
||||
strides = [1,1,strides,strides]
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
dilations = [1,dilations,dilations,1]
|
||||
else:
|
||||
dilations = [1,1,dilations,dilations]
|
||||
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.strides = strides
|
||||
self.padding = padding
|
||||
self.dilations = dilations
|
||||
self.use_bias = use_bias
|
||||
self.use_wscale = use_wscale
|
||||
self.kernel_initializer = kernel_initializer
|
||||
self.bias_initializer = bias_initializer
|
||||
self.trainable = trainable
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = self.kernel_initializer
|
||||
if self.use_wscale:
|
||||
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
|
||||
fan_in = self.kernel_size*self.kernel_size*self.in_ch
|
||||
he_std = gain / np.sqrt(fan_in)
|
||||
self.wscale = tf.constant(he_std, dtype=self.dtype )
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
|
||||
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = nn.initializers.ca()
|
||||
|
||||
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.out_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
|
||||
|
||||
if self.use_bias:
|
||||
bias_initializer = self.bias_initializer
|
||||
if bias_initializer is None:
|
||||
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
|
||||
|
||||
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
|
||||
|
||||
def get_weights(self):
|
||||
weights = [self.weight]
|
||||
if self.use_bias:
|
||||
weights += [self.bias]
|
||||
return weights
|
||||
|
||||
def forward(self, x):
|
||||
weight = self.weight
|
||||
if self.use_wscale:
|
||||
weight = weight * self.wscale
|
||||
|
||||
if self.padding is not None:
|
||||
x = tf.pad (x, self.padding, mode='CONSTANT')
|
||||
|
||||
x = tf.nn.conv2d(x, weight, self.strides, 'VALID', dilations=self.dilations, data_format=nn.data_format)
|
||||
if self.use_bias:
|
||||
if nn.data_format == "NHWC":
|
||||
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
|
||||
else:
|
||||
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
|
||||
x = tf.add(x, bias)
|
||||
return x
|
||||
|
||||
def __str__(self):
|
||||
r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
|
||||
|
||||
return r
|
||||
nn.Conv2D = Conv2D
|
||||
107
core/leras/layers/Conv2DTranspose.py
Normal file
107
core/leras/layers/Conv2DTranspose.py
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
import numpy as np
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class Conv2DTranspose(nn.LayerBase):
|
||||
"""
|
||||
use_wscale enables weight scale (equalized learning rate)
|
||||
if kernel_initializer is None, it will be forced to random_normal
|
||||
"""
|
||||
def __init__(self, in_ch, out_ch, kernel_size, strides=2, padding='SAME', use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
|
||||
if not isinstance(strides, int):
|
||||
raise ValueError ("strides must be an int type")
|
||||
kernel_size = int(kernel_size)
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.kernel_size = kernel_size
|
||||
self.strides = strides
|
||||
self.padding = padding
|
||||
self.use_bias = use_bias
|
||||
self.use_wscale = use_wscale
|
||||
self.kernel_initializer = kernel_initializer
|
||||
self.bias_initializer = bias_initializer
|
||||
self.trainable = trainable
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = self.kernel_initializer
|
||||
if self.use_wscale:
|
||||
gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
|
||||
fan_in = self.kernel_size*self.kernel_size*self.in_ch
|
||||
he_std = gain / np.sqrt(fan_in) # He init
|
||||
self.wscale = tf.constant(he_std, dtype=self.dtype )
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
|
||||
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = nn.initializers.ca()
|
||||
self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.out_ch,self.in_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
|
||||
|
||||
if self.use_bias:
|
||||
bias_initializer = self.bias_initializer
|
||||
if bias_initializer is None:
|
||||
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
|
||||
|
||||
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
|
||||
|
||||
def get_weights(self):
|
||||
weights = [self.weight]
|
||||
if self.use_bias:
|
||||
weights += [self.bias]
|
||||
return weights
|
||||
|
||||
def forward(self, x):
|
||||
shape = x.shape
|
||||
|
||||
if nn.data_format == "NHWC":
|
||||
h,w,c = shape[1], shape[2], shape[3]
|
||||
output_shape = tf.stack ( (tf.shape(x)[0],
|
||||
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
|
||||
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
|
||||
self.out_ch) )
|
||||
|
||||
strides = [1,self.strides,self.strides,1]
|
||||
else:
|
||||
c,h,w = shape[1], shape[2], shape[3]
|
||||
output_shape = tf.stack ( (tf.shape(x)[0],
|
||||
self.out_ch,
|
||||
self.deconv_length(w, self.strides, self.kernel_size, self.padding),
|
||||
self.deconv_length(h, self.strides, self.kernel_size, self.padding),
|
||||
) )
|
||||
strides = [1,1,self.strides,self.strides]
|
||||
weight = self.weight
|
||||
if self.use_wscale:
|
||||
weight = weight * self.wscale
|
||||
|
||||
x = tf.nn.conv2d_transpose(x, weight, output_shape, strides, padding=self.padding, data_format=nn.data_format)
|
||||
|
||||
if self.use_bias:
|
||||
if nn.data_format == "NHWC":
|
||||
bias = tf.reshape (self.bias, (1,1,1,self.out_ch) )
|
||||
else:
|
||||
bias = tf.reshape (self.bias, (1,self.out_ch,1,1) )
|
||||
x = tf.add(x, bias)
|
||||
return x
|
||||
|
||||
def __str__(self):
|
||||
r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
|
||||
|
||||
return r
|
||||
|
||||
def deconv_length(self, dim_size, stride_size, kernel_size, padding):
|
||||
assert padding in {'SAME', 'VALID', 'FULL'}
|
||||
if dim_size is None:
|
||||
return None
|
||||
if padding == 'VALID':
|
||||
dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0)
|
||||
elif padding == 'FULL':
|
||||
dim_size = dim_size * stride_size - (stride_size + kernel_size - 2)
|
||||
elif padding == 'SAME':
|
||||
dim_size = dim_size * stride_size
|
||||
return dim_size
|
||||
nn.Conv2DTranspose = Conv2DTranspose
|
||||
76
core/leras/layers/Dense.py
Normal file
76
core/leras/layers/Dense.py
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
import numpy as np
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class Dense(nn.LayerBase):
|
||||
def __init__(self, in_ch, out_ch, use_bias=True, use_wscale=False, maxout_ch=0, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
|
||||
"""
|
||||
use_wscale enables weight scale (equalized learning rate)
|
||||
if kernel_initializer is None, it will be forced to random_normal
|
||||
|
||||
maxout_ch https://link.springer.com/article/10.1186/s40537-019-0233-0
|
||||
typical 2-4 if you want to enable DenseMaxout behaviour
|
||||
"""
|
||||
self.in_ch = in_ch
|
||||
self.out_ch = out_ch
|
||||
self.use_bias = use_bias
|
||||
self.use_wscale = use_wscale
|
||||
self.maxout_ch = maxout_ch
|
||||
self.kernel_initializer = kernel_initializer
|
||||
self.bias_initializer = bias_initializer
|
||||
self.trainable = trainable
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
|
||||
self.dtype = dtype
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
if self.maxout_ch > 1:
|
||||
weight_shape = (self.in_ch,self.out_ch*self.maxout_ch)
|
||||
else:
|
||||
weight_shape = (self.in_ch,self.out_ch)
|
||||
|
||||
kernel_initializer = self.kernel_initializer
|
||||
|
||||
if self.use_wscale:
|
||||
gain = 1.0
|
||||
fan_in = np.prod( weight_shape[:-1] )
|
||||
he_std = gain / np.sqrt(fan_in) # He init
|
||||
self.wscale = tf.constant(he_std, dtype=self.dtype )
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
|
||||
|
||||
if kernel_initializer is None:
|
||||
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
|
||||
|
||||
self.weight = tf.get_variable("weight", weight_shape, dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
|
||||
|
||||
if self.use_bias:
|
||||
bias_initializer = self.bias_initializer
|
||||
if bias_initializer is None:
|
||||
bias_initializer = tf.initializers.zeros(dtype=self.dtype)
|
||||
self.bias = tf.get_variable("bias", (self.out_ch,), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
|
||||
|
||||
def get_weights(self):
|
||||
weights = [self.weight]
|
||||
if self.use_bias:
|
||||
weights += [self.bias]
|
||||
return weights
|
||||
|
||||
def forward(self, x):
|
||||
weight = self.weight
|
||||
if self.use_wscale:
|
||||
weight = weight * self.wscale
|
||||
|
||||
x = tf.matmul(x, weight)
|
||||
|
||||
if self.maxout_ch > 1:
|
||||
x = tf.reshape (x, (-1, self.out_ch, self.maxout_ch) )
|
||||
x = tf.reduce_max(x, axis=-1)
|
||||
|
||||
if self.use_bias:
|
||||
x = tf.add(x, tf.reshape(self.bias, (1,self.out_ch) ) )
|
||||
|
||||
return x
|
||||
nn.Dense = Dense
|
||||
38
core/leras/layers/FRNorm2D.py
Normal file
38
core/leras/layers/FRNorm2D.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class FRNorm2D(nn.LayerBase):
|
||||
"""
|
||||
Tensorflow implementation of
|
||||
Filter Response Normalization Layer: Eliminating Batch Dependence in theTraining of Deep Neural Networks
|
||||
https://arxiv.org/pdf/1911.09737.pdf
|
||||
"""
|
||||
def __init__(self, in_ch, dtype=None, **kwargs):
|
||||
self.in_ch = in_ch
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
self.dtype = dtype
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
self.weight = tf.get_variable("weight", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.ones() )
|
||||
self.bias = tf.get_variable("bias", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
|
||||
self.eps = tf.get_variable("eps", (1,), dtype=self.dtype, initializer=tf.initializers.constant(1e-6) )
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight, self.bias, self.eps]
|
||||
|
||||
def forward(self, x):
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (1,1,1,self.in_ch)
|
||||
else:
|
||||
shape = (1,self.in_ch,1,1)
|
||||
weight = tf.reshape ( self.weight, shape )
|
||||
bias = tf.reshape ( self.bias , shape )
|
||||
nu2 = tf.reduce_mean(tf.square(x), axis=nn.conv2d_spatial_axes, keepdims=True)
|
||||
x = x * ( 1.0/tf.sqrt(nu2 + tf.abs(self.eps) ) )
|
||||
|
||||
return x*weight + bias
|
||||
nn.FRNorm2D = FRNorm2D
|
||||
40
core/leras/layers/InstanceNorm2D.py
Normal file
40
core/leras/layers/InstanceNorm2D.py
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class InstanceNorm2D(nn.LayerBase):
|
||||
def __init__(self, in_ch, dtype=None, **kwargs):
|
||||
self.in_ch = in_ch
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
self.dtype = dtype
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
|
||||
self.weight = tf.get_variable("weight", (self.in_ch,), dtype=self.dtype, initializer=kernel_initializer )
|
||||
self.bias = tf.get_variable("bias", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight, self.bias]
|
||||
|
||||
def forward(self, x):
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (1,1,1,self.in_ch)
|
||||
else:
|
||||
shape = (1,self.in_ch,1,1)
|
||||
|
||||
weight = tf.reshape ( self.weight , shape )
|
||||
bias = tf.reshape ( self.bias , shape )
|
||||
|
||||
x_mean = tf.reduce_mean(x, axis=nn.conv2d_spatial_axes, keepdims=True )
|
||||
x_std = tf.math.reduce_std(x, axis=nn.conv2d_spatial_axes, keepdims=True ) + 1e-5
|
||||
|
||||
x = (x - x_mean) / x_std
|
||||
x *= weight
|
||||
x += bias
|
||||
|
||||
return x
|
||||
|
||||
nn.InstanceNorm2D = InstanceNorm2D
|
||||
16
core/leras/layers/LayerBase.py
Normal file
16
core/leras/layers/LayerBase.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class LayerBase(nn.Saveable):
|
||||
#override
|
||||
def build_weights(self):
|
||||
pass
|
||||
|
||||
#override
|
||||
def forward(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self.forward(*args, **kwargs)
|
||||
|
||||
nn.LayerBase = LayerBase
|
||||
103
core/leras/layers/Saveable.py
Normal file
103
core/leras/layers/Saveable.py
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
import pickle
|
||||
from pathlib import Path
|
||||
from core import pathex
|
||||
import numpy as np
|
||||
|
||||
from core.leras import nn
|
||||
|
||||
tf = nn.tf
|
||||
|
||||
class Saveable():
|
||||
def __init__(self, name=None):
|
||||
self.name = name
|
||||
|
||||
#override
|
||||
def get_weights(self):
|
||||
#return tf tensors that should be initialized/loaded/saved
|
||||
return []
|
||||
|
||||
#override
|
||||
def get_weights_np(self):
|
||||
weights = self.get_weights()
|
||||
if len(weights) == 0:
|
||||
return []
|
||||
return nn.tf_sess.run (weights)
|
||||
|
||||
def set_weights(self, new_weights):
|
||||
weights = self.get_weights()
|
||||
if len(weights) != len(new_weights):
|
||||
raise ValueError ('len of lists mismatch')
|
||||
|
||||
tuples = []
|
||||
for w, new_w in zip(weights, new_weights):
|
||||
|
||||
if len(w.shape) != new_w.shape:
|
||||
new_w = new_w.reshape(w.shape)
|
||||
|
||||
tuples.append ( (w, new_w) )
|
||||
|
||||
nn.batch_set_value (tuples)
|
||||
|
||||
def save_weights(self, filename, force_dtype=None):
|
||||
d = {}
|
||||
weights = self.get_weights()
|
||||
|
||||
if self.name is None:
|
||||
raise Exception("name must be defined.")
|
||||
|
||||
name = self.name
|
||||
for w, w_val in zip(weights, nn.tf_sess.run (weights)):
|
||||
w_name_split = w.name.split('/', 1)
|
||||
if name != w_name_split[0]:
|
||||
raise Exception("weight first name != Saveable.name")
|
||||
|
||||
if force_dtype is not None:
|
||||
w_val = w_val.astype(force_dtype)
|
||||
|
||||
d[ w_name_split[1] ] = w_val
|
||||
|
||||
d_dumped = pickle.dumps (d, 4)
|
||||
pathex.write_bytes_safe ( Path(filename), d_dumped )
|
||||
|
||||
def load_weights(self, filename):
|
||||
"""
|
||||
returns True if file exists
|
||||
"""
|
||||
filepath = Path(filename)
|
||||
if filepath.exists():
|
||||
result = True
|
||||
d_dumped = filepath.read_bytes()
|
||||
d = pickle.loads(d_dumped)
|
||||
else:
|
||||
return False
|
||||
|
||||
weights = self.get_weights()
|
||||
|
||||
if self.name is None:
|
||||
raise Exception("name must be defined.")
|
||||
|
||||
tuples = []
|
||||
for w in weights:
|
||||
w_name_split = w.name.split('/')
|
||||
if self.name != w_name_split[0]:
|
||||
raise Exception("weight first name != Saveable.name")
|
||||
|
||||
sub_w_name = "/".join(w_name_split[1:])
|
||||
|
||||
w_val = d.get(sub_w_name, None)
|
||||
|
||||
if w_val is None:
|
||||
#io.log_err(f"Weight {w.name} was not loaded from file {filename}")
|
||||
tuples.append ( (w, w.initializer) )
|
||||
else:
|
||||
w_val = np.reshape( w_val, w.shape.as_list() )
|
||||
tuples.append ( (w, w_val) )
|
||||
|
||||
nn.batch_set_value(tuples)
|
||||
|
||||
return True
|
||||
|
||||
def init_weights(self):
|
||||
nn.init_weights(self.get_weights())
|
||||
|
||||
nn.Saveable = Saveable
|
||||
33
core/leras/layers/TLU.py
Normal file
33
core/leras/layers/TLU.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class TLU(nn.LayerBase):
|
||||
"""
|
||||
Tensorflow implementation of
|
||||
Filter Response Normalization Layer: Eliminating Batch Dependence in theTraining of Deep Neural Networks
|
||||
https://arxiv.org/pdf/1911.09737.pdf
|
||||
"""
|
||||
def __init__(self, in_ch, dtype=None, **kwargs):
|
||||
self.in_ch = in_ch
|
||||
|
||||
if dtype is None:
|
||||
dtype = nn.floatx
|
||||
self.dtype = dtype
|
||||
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
self.tau = tf.get_variable("tau", (self.in_ch,), dtype=self.dtype, initializer=tf.initializers.zeros() )
|
||||
|
||||
def get_weights(self):
|
||||
return [self.tau]
|
||||
|
||||
def forward(self, x):
|
||||
if nn.data_format == "NHWC":
|
||||
shape = (1,1,1,self.in_ch)
|
||||
else:
|
||||
shape = (1,self.in_ch,1,1)
|
||||
|
||||
tau = tf.reshape ( self.tau, shape )
|
||||
return tf.math.maximum(x, tau)
|
||||
nn.TLU = TLU
|
||||
12
core/leras/layers/__init__.py
Normal file
12
core/leras/layers/__init__.py
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
from .Saveable import *
|
||||
from .LayerBase import *
|
||||
|
||||
from .Conv2D import *
|
||||
from .Conv2DTranspose import *
|
||||
from .Dense import *
|
||||
from .BlurPool import *
|
||||
|
||||
from .BatchNorm2D import *
|
||||
from .FRNorm2D import *
|
||||
|
||||
from .TLU import *
|
||||
|
|
@ -1,367 +0,0 @@
|
|||
import types
|
||||
import numpy as np
|
||||
from core.interact import interact as io
|
||||
|
||||
def initialize_models(nn):
|
||||
tf = nn.tf
|
||||
|
||||
class ModelBase(nn.Saveable):
|
||||
def __init__(self, *args, name=None, **kwargs):
|
||||
super().__init__(name=name)
|
||||
self.layers = []
|
||||
self.layers_by_name = {}
|
||||
self.built = False
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
self.run_placeholders = None
|
||||
|
||||
def _build_sub(self, layer, name):
|
||||
if isinstance (layer, list):
|
||||
for i,sublayer in enumerate(layer):
|
||||
self._build_sub(sublayer, f"{name}_{i}")
|
||||
elif isinstance (layer, nn.LayerBase) or \
|
||||
isinstance (layer, ModelBase):
|
||||
|
||||
if layer.name is None:
|
||||
layer.name = name
|
||||
|
||||
if isinstance (layer, nn.LayerBase):
|
||||
with tf.variable_scope(layer.name):
|
||||
layer.build_weights()
|
||||
elif isinstance (layer, ModelBase):
|
||||
layer.build()
|
||||
|
||||
self.layers.append (layer)
|
||||
self.layers_by_name[layer.name] = layer
|
||||
|
||||
def xor_list(self, lst1, lst2):
|
||||
return [value for value in lst1+lst2 if (value not in lst1) or (value not in lst2) ]
|
||||
|
||||
def build(self):
|
||||
with tf.variable_scope(self.name):
|
||||
|
||||
current_vars = []
|
||||
generator = None
|
||||
while True:
|
||||
|
||||
if generator is None:
|
||||
generator = self.on_build(*self.args, **self.kwargs)
|
||||
if not isinstance(generator, types.GeneratorType):
|
||||
generator = None
|
||||
|
||||
if generator is not None:
|
||||
try:
|
||||
next(generator)
|
||||
except StopIteration:
|
||||
generator = None
|
||||
|
||||
v = vars(self)
|
||||
new_vars = self.xor_list (current_vars, list(v.keys()) )
|
||||
|
||||
for name in new_vars:
|
||||
self._build_sub(v[name],name)
|
||||
|
||||
current_vars += new_vars
|
||||
|
||||
if generator is None:
|
||||
break
|
||||
|
||||
self.built = True
|
||||
|
||||
#override
|
||||
def get_weights(self):
|
||||
if not self.built:
|
||||
self.build()
|
||||
|
||||
weights = []
|
||||
for layer in self.layers:
|
||||
weights += layer.get_weights()
|
||||
return weights
|
||||
|
||||
def get_layer_by_name(self, name):
|
||||
return self.layers_by_name.get(name, None)
|
||||
|
||||
def get_layers(self):
|
||||
if not self.built:
|
||||
self.build()
|
||||
layers = []
|
||||
for layer in self.layers:
|
||||
if isinstance (layer, nn.LayerBase):
|
||||
layers.append(layer)
|
||||
else:
|
||||
layers += layer.get_layers()
|
||||
return layers
|
||||
|
||||
#override
|
||||
def on_build(self, *args, **kwargs):
|
||||
"""
|
||||
init model layers here
|
||||
|
||||
return 'yield' if build is not finished
|
||||
therefore dependency models will be initialized
|
||||
"""
|
||||
pass
|
||||
|
||||
#override
|
||||
def forward(self, *args, **kwargs):
|
||||
#flow layers/models/tensors here
|
||||
pass
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
if not self.built:
|
||||
self.build()
|
||||
|
||||
return self.forward(*args, **kwargs)
|
||||
|
||||
def compute_output_shape(self, shapes):
|
||||
if not self.built:
|
||||
self.build()
|
||||
|
||||
not_list = False
|
||||
if not isinstance(shapes, list):
|
||||
not_list = True
|
||||
shapes = [shapes]
|
||||
|
||||
with tf.device('/CPU:0'):
|
||||
# CPU tensors will not impact any performance, only slightly RAM "leakage"
|
||||
phs = []
|
||||
for dtype,sh in shapes:
|
||||
phs += [ tf.placeholder(dtype, sh) ]
|
||||
|
||||
result = self.__call__(phs[0] if not_list else phs)
|
||||
|
||||
if not isinstance(result, list):
|
||||
result = [result]
|
||||
|
||||
result_shapes = []
|
||||
|
||||
for t in result:
|
||||
result_shapes += [ t.shape.as_list() ]
|
||||
|
||||
return result_shapes[0] if not_list else result_shapes
|
||||
|
||||
def compute_output_channels(self, shapes):
|
||||
shape = self.compute_output_shape(shapes)
|
||||
shape_len = len(shape)
|
||||
|
||||
if shape_len == 4:
|
||||
if nn.data_format == "NCHW":
|
||||
return shape[1]
|
||||
return shape[-1]
|
||||
|
||||
def build_for_run(self, shapes_list):
|
||||
if not isinstance(shapes_list, list):
|
||||
raise ValueError("shapes_list must be a list.")
|
||||
|
||||
self.run_placeholders = []
|
||||
for dtype,sh in shapes_list:
|
||||
self.run_placeholders.append ( tf.placeholder(dtype, sh) )
|
||||
|
||||
self.run_output = self.__call__(self.run_placeholders)
|
||||
|
||||
def run (self, inputs):
|
||||
if self.run_placeholders is None:
|
||||
raise Exception ("Model didn't build for run.")
|
||||
|
||||
if len(inputs) != len(self.run_placeholders):
|
||||
raise ValueError("len(inputs) != self.run_placeholders")
|
||||
|
||||
feed_dict = {}
|
||||
for ph, inp in zip(self.run_placeholders, inputs):
|
||||
feed_dict[ph] = inp
|
||||
|
||||
return nn.tf_sess.run ( self.run_output, feed_dict=feed_dict)
|
||||
|
||||
def summary(self):
|
||||
layers = self.get_layers()
|
||||
layers_names = []
|
||||
layers_params = []
|
||||
|
||||
max_len_str = 0
|
||||
max_len_param_str = 0
|
||||
delim_str = "-"
|
||||
|
||||
total_params = 0
|
||||
|
||||
#Get layers names and str lenght for delim
|
||||
for l in layers:
|
||||
if len(str(l))>max_len_str:
|
||||
max_len_str = len(str(l))
|
||||
layers_names+=[str(l).capitalize()]
|
||||
|
||||
#Get params for each layer
|
||||
layers_params = [ int(np.sum(np.prod(w.shape) for w in l.get_weights())) for l in layers ]
|
||||
total_params = np.sum(layers_params)
|
||||
|
||||
#Get str lenght for delim
|
||||
for p in layers_params:
|
||||
if len(str(p))>max_len_param_str:
|
||||
max_len_param_str=len(str(p))
|
||||
|
||||
#Set delim
|
||||
for i in range(max_len_str+max_len_param_str+3):
|
||||
delim_str += "-"
|
||||
|
||||
output = "\n"+delim_str+"\n"
|
||||
|
||||
#Format model name str
|
||||
model_name_str = "| "+self.name.capitalize()
|
||||
len_model_name_str = len(model_name_str)
|
||||
for i in range(len(delim_str)-len_model_name_str):
|
||||
model_name_str+= " " if i!=(len(delim_str)-len_model_name_str-2) else " |"
|
||||
|
||||
output += model_name_str +"\n"
|
||||
output += delim_str +"\n"
|
||||
|
||||
|
||||
#Format layers table
|
||||
for i in range(len(layers_names)):
|
||||
output += delim_str +"\n"
|
||||
|
||||
l_name = layers_names[i]
|
||||
l_param = str(layers_params[i])
|
||||
l_param_str = ""
|
||||
if len(l_name)<=max_len_str:
|
||||
for i in range(max_len_str - len(l_name)):
|
||||
l_name+= " "
|
||||
|
||||
if len(l_param)<=max_len_param_str:
|
||||
for i in range(max_len_param_str - len(l_param)):
|
||||
l_param_str+= " "
|
||||
|
||||
l_param_str += l_param
|
||||
|
||||
|
||||
output +="| "+l_name+"|"+l_param_str+"| \n"
|
||||
|
||||
output += delim_str +"\n"
|
||||
|
||||
#Format sum of params
|
||||
total_params_str = "| Total params count: "+str(total_params)
|
||||
len_total_params_str = len(total_params_str)
|
||||
for i in range(len(delim_str)-len_total_params_str):
|
||||
total_params_str+= " " if i!=(len(delim_str)-len_total_params_str-2) else " |"
|
||||
|
||||
output += total_params_str +"\n"
|
||||
output += delim_str +"\n"
|
||||
|
||||
io.log_info(output)
|
||||
|
||||
nn.ModelBase = ModelBase
|
||||
|
||||
class PatchDiscriminator(nn.ModelBase):
|
||||
def on_build(self, patch_size, in_ch, base_ch=None, conv_kernel_initializer=None):
|
||||
suggested_base_ch, kernels_strides = patch_discriminator_kernels[patch_size]
|
||||
|
||||
if base_ch is None:
|
||||
base_ch = suggested_base_ch
|
||||
|
||||
prev_ch = in_ch
|
||||
self.convs = []
|
||||
for i, (kernel_size, strides) in enumerate(kernels_strides):
|
||||
cur_ch = base_ch * min( (2**i), 8 )
|
||||
|
||||
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=kernel_size, strides=strides, padding='SAME', kernel_initializer=conv_kernel_initializer) )
|
||||
prev_ch = cur_ch
|
||||
|
||||
self.out_conv = nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
for conv in self.convs:
|
||||
x = tf.nn.leaky_relu( conv(x), 0.1 )
|
||||
return self.out_conv(x)
|
||||
|
||||
nn.PatchDiscriminator = PatchDiscriminator
|
||||
|
||||
class IllumDiscriminator(nn.ModelBase):
|
||||
def on_build(self, patch_size, in_ch, base_ch=None, conv_kernel_initializer=None):
|
||||
suggested_base_ch, kernels_strides = patch_discriminator_kernels[patch_size]
|
||||
if base_ch is None:
|
||||
base_ch = suggested_base_ch
|
||||
|
||||
prev_ch = in_ch
|
||||
self.convs = []
|
||||
for i, (kernel_size, strides) in enumerate(kernels_strides):
|
||||
cur_ch = base_ch * min( (2**i), 8 )
|
||||
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=kernel_size, strides=strides, padding='SAME', kernel_initializer=conv_kernel_initializer) )
|
||||
prev_ch = cur_ch
|
||||
|
||||
self.out1 = nn.Conv2D( 1, 1024, kernel_size=1, strides=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
self.out2 = nn.Conv2D( 1024, 1, kernel_size=1, strides=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
for conv in self.convs:
|
||||
x = tf.nn.leaky_relu( conv(x), 0.1 )
|
||||
|
||||
x = tf.reduce_mean(x, axis=nn.conv2d_ch_axis, keep_dims=True)
|
||||
|
||||
x = self.out1(x)
|
||||
x = tf.nn.leaky_relu(x, 0.1 )
|
||||
x = self.out2(x)
|
||||
|
||||
return x
|
||||
|
||||
nn.IllumDiscriminator = IllumDiscriminator
|
||||
|
||||
class CodeDiscriminator(nn.ModelBase):
|
||||
def on_build(self, in_ch, code_res, ch=256, conv_kernel_initializer=None):
|
||||
if conv_kernel_initializer is None:
|
||||
conv_kernel_initializer = nn.initializers.ca()
|
||||
|
||||
n_downscales = 1 + code_res // 8
|
||||
|
||||
self.convs = []
|
||||
prev_ch = in_ch
|
||||
for i in range(n_downscales):
|
||||
cur_ch = ch * min( (2**i), 8 )
|
||||
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=4 if i == 0 else 3, strides=2, padding='SAME', kernel_initializer=conv_kernel_initializer) )
|
||||
prev_ch = cur_ch
|
||||
|
||||
self.out_conv = nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
for conv in self.convs:
|
||||
x = tf.nn.leaky_relu( conv(x), 0.1 )
|
||||
return self.out_conv(x)
|
||||
nn.CodeDiscriminator = CodeDiscriminator
|
||||
|
||||
patch_discriminator_kernels = \
|
||||
{ 1 : (512, [ [1,1] ]),
|
||||
2 : (512, [ [2,1] ]),
|
||||
3 : (512, [ [2,1], [2,1] ]),
|
||||
4 : (512, [ [2,2], [2,2] ]),
|
||||
5 : (512, [ [3,2], [2,2] ]),
|
||||
6 : (512, [ [4,2], [2,2] ]),
|
||||
7 : (512, [ [3,2], [3,2] ]),
|
||||
8 : (512, [ [4,2], [3,2] ]),
|
||||
9 : (512, [ [3,2], [4,2] ]),
|
||||
10 : (512, [ [4,2], [4,2] ]),
|
||||
11 : (512, [ [3,2], [3,2], [2,1] ]),
|
||||
12 : (512, [ [4,2], [3,2], [2,1] ]),
|
||||
13 : (512, [ [3,2], [4,2], [2,1] ]),
|
||||
14 : (512, [ [4,2], [4,2], [2,1] ]),
|
||||
15 : (512, [ [3,2], [3,2], [3,1] ]),
|
||||
16 : (512, [ [4,2], [3,2], [3,1] ]),
|
||||
17 : (512, [ [3,2], [4,2], [3,1] ]),
|
||||
18 : (512, [ [4,2], [4,2], [3,1] ]),
|
||||
19 : (512, [ [3,2], [3,2], [4,1] ]),
|
||||
20 : (512, [ [4,2], [3,2], [4,1] ]),
|
||||
21 : (512, [ [3,2], [4,2], [4,1] ]),
|
||||
22 : (512, [ [4,2], [4,2], [4,1] ]),
|
||||
23 : (256, [ [3,2], [3,2], [3,2], [2,1] ]),
|
||||
24 : (256, [ [4,2], [3,2], [3,2], [2,1] ]),
|
||||
25 : (256, [ [3,2], [4,2], [3,2], [2,1] ]),
|
||||
26 : (256, [ [4,2], [4,2], [3,2], [2,1] ]),
|
||||
27 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
|
||||
28 : (256, [ [4,2], [3,2], [4,2], [2,1] ]),
|
||||
29 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
|
||||
30 : (256, [ [4,2], [4,2], [4,2], [2,1] ]),
|
||||
31 : (256, [ [3,2], [3,2], [3,2], [3,1] ]),
|
||||
32 : (256, [ [4,2], [3,2], [3,2], [3,1] ]),
|
||||
33 : (256, [ [3,2], [4,2], [3,2], [3,1] ]),
|
||||
34 : (256, [ [4,2], [4,2], [3,2], [3,1] ]),
|
||||
35 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
|
||||
36 : (256, [ [4,2], [3,2], [4,2], [3,1] ]),
|
||||
37 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
|
||||
38 : (256, [ [4,2], [4,2], [4,2], [3,1] ]),
|
||||
}
|
||||
22
core/leras/models/CodeDiscriminator.py
Normal file
22
core/leras/models/CodeDiscriminator.py
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class CodeDiscriminator(nn.ModelBase):
|
||||
def on_build(self, in_ch, code_res, ch=256, conv_kernel_initializer=None):
|
||||
n_downscales = 1 + code_res // 8
|
||||
|
||||
self.convs = []
|
||||
prev_ch = in_ch
|
||||
for i in range(n_downscales):
|
||||
cur_ch = ch * min( (2**i), 8 )
|
||||
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=4 if i == 0 else 3, strides=2, padding='SAME', kernel_initializer=conv_kernel_initializer) )
|
||||
prev_ch = cur_ch
|
||||
|
||||
self.out_conv = nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
for conv in self.convs:
|
||||
x = tf.nn.leaky_relu( conv(x), 0.1 )
|
||||
return self.out_conv(x)
|
||||
|
||||
nn.CodeDiscriminator = CodeDiscriminator
|
||||
249
core/leras/models/ModelBase.py
Normal file
249
core/leras/models/ModelBase.py
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
import types
|
||||
import numpy as np
|
||||
from core.interact import interact as io
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class ModelBase(nn.Saveable):
|
||||
def __init__(self, *args, name=None, **kwargs):
|
||||
super().__init__(name=name)
|
||||
self.layers = []
|
||||
self.layers_by_name = {}
|
||||
self.built = False
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
self.run_placeholders = None
|
||||
|
||||
def _build_sub(self, layer, name):
|
||||
if isinstance (layer, list):
|
||||
for i,sublayer in enumerate(layer):
|
||||
self._build_sub(sublayer, f"{name}_{i}")
|
||||
elif isinstance (layer, nn.LayerBase) or \
|
||||
isinstance (layer, ModelBase):
|
||||
|
||||
if layer.name is None:
|
||||
layer.name = name
|
||||
|
||||
if isinstance (layer, nn.LayerBase):
|
||||
with tf.variable_scope(layer.name):
|
||||
layer.build_weights()
|
||||
elif isinstance (layer, ModelBase):
|
||||
layer.build()
|
||||
|
||||
self.layers.append (layer)
|
||||
self.layers_by_name[layer.name] = layer
|
||||
|
||||
def xor_list(self, lst1, lst2):
|
||||
return [value for value in lst1+lst2 if (value not in lst1) or (value not in lst2) ]
|
||||
|
||||
def build(self):
|
||||
with tf.variable_scope(self.name):
|
||||
|
||||
current_vars = []
|
||||
generator = None
|
||||
while True:
|
||||
|
||||
if generator is None:
|
||||
generator = self.on_build(*self.args, **self.kwargs)
|
||||
if not isinstance(generator, types.GeneratorType):
|
||||
generator = None
|
||||
|
||||
if generator is not None:
|
||||
try:
|
||||
next(generator)
|
||||
except StopIteration:
|
||||
generator = None
|
||||
|
||||
v = vars(self)
|
||||
new_vars = self.xor_list (current_vars, list(v.keys()) )
|
||||
|
||||
for name in new_vars:
|
||||
self._build_sub(v[name],name)
|
||||
|
||||
current_vars += new_vars
|
||||
|
||||
if generator is None:
|
||||
break
|
||||
|
||||
self.built = True
|
||||
|
||||
#override
|
||||
def get_weights(self):
|
||||
if not self.built:
|
||||
self.build()
|
||||
|
||||
weights = []
|
||||
for layer in self.layers:
|
||||
weights += layer.get_weights()
|
||||
return weights
|
||||
|
||||
def get_layer_by_name(self, name):
|
||||
return self.layers_by_name.get(name, None)
|
||||
|
||||
def get_layers(self):
|
||||
if not self.built:
|
||||
self.build()
|
||||
layers = []
|
||||
for layer in self.layers:
|
||||
if isinstance (layer, nn.LayerBase):
|
||||
layers.append(layer)
|
||||
else:
|
||||
layers += layer.get_layers()
|
||||
return layers
|
||||
|
||||
#override
|
||||
def on_build(self, *args, **kwargs):
|
||||
"""
|
||||
init model layers here
|
||||
|
||||
return 'yield' if build is not finished
|
||||
therefore dependency models will be initialized
|
||||
"""
|
||||
pass
|
||||
|
||||
#override
|
||||
def forward(self, *args, **kwargs):
|
||||
#flow layers/models/tensors here
|
||||
pass
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
if not self.built:
|
||||
self.build()
|
||||
|
||||
return self.forward(*args, **kwargs)
|
||||
|
||||
def compute_output_shape(self, shapes):
|
||||
if not self.built:
|
||||
self.build()
|
||||
|
||||
not_list = False
|
||||
if not isinstance(shapes, list):
|
||||
not_list = True
|
||||
shapes = [shapes]
|
||||
|
||||
with tf.device('/CPU:0'):
|
||||
# CPU tensors will not impact any performance, only slightly RAM "leakage"
|
||||
phs = []
|
||||
for dtype,sh in shapes:
|
||||
phs += [ tf.placeholder(dtype, sh) ]
|
||||
|
||||
result = self.__call__(phs[0] if not_list else phs)
|
||||
|
||||
if not isinstance(result, list):
|
||||
result = [result]
|
||||
|
||||
result_shapes = []
|
||||
|
||||
for t in result:
|
||||
result_shapes += [ t.shape.as_list() ]
|
||||
|
||||
return result_shapes[0] if not_list else result_shapes
|
||||
|
||||
def compute_output_channels(self, shapes):
|
||||
shape = self.compute_output_shape(shapes)
|
||||
shape_len = len(shape)
|
||||
|
||||
if shape_len == 4:
|
||||
if nn.data_format == "NCHW":
|
||||
return shape[1]
|
||||
return shape[-1]
|
||||
|
||||
def build_for_run(self, shapes_list):
|
||||
if not isinstance(shapes_list, list):
|
||||
raise ValueError("shapes_list must be a list.")
|
||||
|
||||
self.run_placeholders = []
|
||||
for dtype,sh in shapes_list:
|
||||
self.run_placeholders.append ( tf.placeholder(dtype, sh) )
|
||||
|
||||
self.run_output = self.__call__(self.run_placeholders)
|
||||
|
||||
def run (self, inputs):
|
||||
if self.run_placeholders is None:
|
||||
raise Exception ("Model didn't build for run.")
|
||||
|
||||
if len(inputs) != len(self.run_placeholders):
|
||||
raise ValueError("len(inputs) != self.run_placeholders")
|
||||
|
||||
feed_dict = {}
|
||||
for ph, inp in zip(self.run_placeholders, inputs):
|
||||
feed_dict[ph] = inp
|
||||
|
||||
return nn.tf_sess.run ( self.run_output, feed_dict=feed_dict)
|
||||
|
||||
def summary(self):
|
||||
layers = self.get_layers()
|
||||
layers_names = []
|
||||
layers_params = []
|
||||
|
||||
max_len_str = 0
|
||||
max_len_param_str = 0
|
||||
delim_str = "-"
|
||||
|
||||
total_params = 0
|
||||
|
||||
#Get layers names and str lenght for delim
|
||||
for l in layers:
|
||||
if len(str(l))>max_len_str:
|
||||
max_len_str = len(str(l))
|
||||
layers_names+=[str(l).capitalize()]
|
||||
|
||||
#Get params for each layer
|
||||
layers_params = [ int(np.sum(np.prod(w.shape) for w in l.get_weights())) for l in layers ]
|
||||
total_params = np.sum(layers_params)
|
||||
|
||||
#Get str lenght for delim
|
||||
for p in layers_params:
|
||||
if len(str(p))>max_len_param_str:
|
||||
max_len_param_str=len(str(p))
|
||||
|
||||
#Set delim
|
||||
for i in range(max_len_str+max_len_param_str+3):
|
||||
delim_str += "-"
|
||||
|
||||
output = "\n"+delim_str+"\n"
|
||||
|
||||
#Format model name str
|
||||
model_name_str = "| "+self.name.capitalize()
|
||||
len_model_name_str = len(model_name_str)
|
||||
for i in range(len(delim_str)-len_model_name_str):
|
||||
model_name_str+= " " if i!=(len(delim_str)-len_model_name_str-2) else " |"
|
||||
|
||||
output += model_name_str +"\n"
|
||||
output += delim_str +"\n"
|
||||
|
||||
|
||||
#Format layers table
|
||||
for i in range(len(layers_names)):
|
||||
output += delim_str +"\n"
|
||||
|
||||
l_name = layers_names[i]
|
||||
l_param = str(layers_params[i])
|
||||
l_param_str = ""
|
||||
if len(l_name)<=max_len_str:
|
||||
for i in range(max_len_str - len(l_name)):
|
||||
l_name+= " "
|
||||
|
||||
if len(l_param)<=max_len_param_str:
|
||||
for i in range(max_len_param_str - len(l_param)):
|
||||
l_param_str+= " "
|
||||
|
||||
l_param_str += l_param
|
||||
|
||||
|
||||
output +="| "+l_name+"|"+l_param_str+"| \n"
|
||||
|
||||
output += delim_str +"\n"
|
||||
|
||||
#Format sum of params
|
||||
total_params_str = "| Total params count: "+str(total_params)
|
||||
len_total_params_str = len(total_params_str)
|
||||
for i in range(len(delim_str)-len_total_params_str):
|
||||
total_params_str+= " " if i!=(len(delim_str)-len_total_params_str-2) else " |"
|
||||
|
||||
output += total_params_str +"\n"
|
||||
output += delim_str +"\n"
|
||||
|
||||
io.log_info(output)
|
||||
|
||||
nn.ModelBase = ModelBase
|
||||
69
core/leras/models/PatchDiscriminator.py
Normal file
69
core/leras/models/PatchDiscriminator.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
|
||||
patch_discriminator_kernels = \
|
||||
{ 1 : (512, [ [1,1] ]),
|
||||
2 : (512, [ [2,1] ]),
|
||||
3 : (512, [ [2,1], [2,1] ]),
|
||||
4 : (512, [ [2,2], [2,2] ]),
|
||||
5 : (512, [ [3,2], [2,2] ]),
|
||||
6 : (512, [ [4,2], [2,2] ]),
|
||||
7 : (512, [ [3,2], [3,2] ]),
|
||||
8 : (512, [ [4,2], [3,2] ]),
|
||||
9 : (512, [ [3,2], [4,2] ]),
|
||||
10 : (512, [ [4,2], [4,2] ]),
|
||||
11 : (512, [ [3,2], [3,2], [2,1] ]),
|
||||
12 : (512, [ [4,2], [3,2], [2,1] ]),
|
||||
13 : (512, [ [3,2], [4,2], [2,1] ]),
|
||||
14 : (512, [ [4,2], [4,2], [2,1] ]),
|
||||
15 : (512, [ [3,2], [3,2], [3,1] ]),
|
||||
16 : (512, [ [4,2], [3,2], [3,1] ]),
|
||||
17 : (512, [ [3,2], [4,2], [3,1] ]),
|
||||
18 : (512, [ [4,2], [4,2], [3,1] ]),
|
||||
19 : (512, [ [3,2], [3,2], [4,1] ]),
|
||||
20 : (512, [ [4,2], [3,2], [4,1] ]),
|
||||
21 : (512, [ [3,2], [4,2], [4,1] ]),
|
||||
22 : (512, [ [4,2], [4,2], [4,1] ]),
|
||||
23 : (256, [ [3,2], [3,2], [3,2], [2,1] ]),
|
||||
24 : (256, [ [4,2], [3,2], [3,2], [2,1] ]),
|
||||
25 : (256, [ [3,2], [4,2], [3,2], [2,1] ]),
|
||||
26 : (256, [ [4,2], [4,2], [3,2], [2,1] ]),
|
||||
27 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
|
||||
28 : (256, [ [4,2], [3,2], [4,2], [2,1] ]),
|
||||
29 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
|
||||
30 : (256, [ [4,2], [4,2], [4,2], [2,1] ]),
|
||||
31 : (256, [ [3,2], [3,2], [3,2], [3,1] ]),
|
||||
32 : (256, [ [4,2], [3,2], [3,2], [3,1] ]),
|
||||
33 : (256, [ [3,2], [4,2], [3,2], [3,1] ]),
|
||||
34 : (256, [ [4,2], [4,2], [3,2], [3,1] ]),
|
||||
35 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
|
||||
36 : (256, [ [4,2], [3,2], [4,2], [3,1] ]),
|
||||
37 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
|
||||
38 : (256, [ [4,2], [4,2], [4,2], [3,1] ]),
|
||||
}
|
||||
|
||||
|
||||
class PatchDiscriminator(nn.ModelBase):
|
||||
def on_build(self, patch_size, in_ch, base_ch=None, conv_kernel_initializer=None):
|
||||
suggested_base_ch, kernels_strides = patch_discriminator_kernels[patch_size]
|
||||
|
||||
if base_ch is None:
|
||||
base_ch = suggested_base_ch
|
||||
|
||||
prev_ch = in_ch
|
||||
self.convs = []
|
||||
for i, (kernel_size, strides) in enumerate(kernels_strides):
|
||||
cur_ch = base_ch * min( (2**i), 8 )
|
||||
|
||||
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=kernel_size, strides=strides, padding='SAME', kernel_initializer=conv_kernel_initializer) )
|
||||
prev_ch = cur_ch
|
||||
|
||||
self.out_conv = nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
|
||||
|
||||
def forward(self, x):
|
||||
for conv in self.convs:
|
||||
x = tf.nn.leaky_relu( conv(x), 0.1 )
|
||||
return self.out_conv(x)
|
||||
|
||||
nn.PatchDiscriminator = PatchDiscriminator
|
||||
4
core/leras/models/__init__.py
Normal file
4
core/leras/models/__init__.py
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
from .ModelBase import *
|
||||
from .PatchDiscriminator import *
|
||||
from .CodeDiscriminator import *
|
||||
from .Ternaus import *
|
||||
116
core/leras/nn.py
116
core/leras/nn.py
|
|
@ -11,7 +11,7 @@ Provides:
|
|||
+ convenient and understandable logic
|
||||
|
||||
Reasons why we cannot import tensorflow or any tensorflow.sub modules right here:
|
||||
1) change env variables based on DeviceConfig before import tensorflow
|
||||
1) program is changing env variables based on DeviceConfig before import tensorflow
|
||||
2) multiprocesses will import tensorflow every spawn
|
||||
|
||||
NCHW speed up training for 10-20%.
|
||||
|
|
@ -19,12 +19,11 @@ NCHW speed up training for 10-20%.
|
|||
|
||||
import os
|
||||
import sys
|
||||
import warnings
|
||||
warnings.simplefilter(action='ignore', category=FutureWarning)
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from core.interact import interact as io
|
||||
|
||||
from .device import Devices
|
||||
|
||||
|
||||
|
|
@ -40,57 +39,7 @@ class nn():
|
|||
conv2d_ch_axis = None
|
||||
conv2d_spatial_axes = None
|
||||
|
||||
tf_floatx = None
|
||||
np_floatx = None
|
||||
|
||||
# Tensor ops
|
||||
tf_get_value = None
|
||||
tf_batch_set_value = None
|
||||
tf_init_weights = None
|
||||
tf_gradients = None
|
||||
tf_average_gv_list = None
|
||||
tf_average_tensor_list = None
|
||||
tf_concat = None
|
||||
tf_gelu = None
|
||||
tf_upsample2d = None
|
||||
tf_resize2d_bilinear = None
|
||||
tf_flatten = None
|
||||
tf_max_pool = None
|
||||
tf_reshape_4D = None
|
||||
tf_random_binomial = None
|
||||
tf_gaussian_blur = None
|
||||
tf_style_loss = None
|
||||
tf_dssim = None
|
||||
tf_space_to_depth = None
|
||||
tf_depth_to_space = None
|
||||
|
||||
# Layers
|
||||
Saveable = None
|
||||
LayerBase = None
|
||||
ModelBase = None
|
||||
Conv2D = None
|
||||
Conv2DTranspose = None
|
||||
BlurPool = None
|
||||
Dense = None
|
||||
InstanceNorm2D = None
|
||||
BatchNorm2D = None
|
||||
AdaIN = None
|
||||
|
||||
# Initializers
|
||||
initializers = None
|
||||
|
||||
# Optimizers
|
||||
TFBaseOptimizer = None
|
||||
TFRMSpropOptimizer = None
|
||||
|
||||
# Models
|
||||
PatchDiscriminator = None
|
||||
IllumDiscriminator = None
|
||||
CodeDiscriminator = None
|
||||
|
||||
# Arhis
|
||||
get_ae_models = None
|
||||
get_ae_models_chervonij = None
|
||||
floatx = None
|
||||
|
||||
@staticmethod
|
||||
def initialize(device_config=None, floatx="float32", data_format="NHWC"):
|
||||
|
|
@ -98,15 +47,17 @@ class nn():
|
|||
if nn.tf is None:
|
||||
if device_config is None:
|
||||
device_config = nn.getCurrentDeviceConfig()
|
||||
else:
|
||||
nn.setCurrentDeviceConfig(device_config)
|
||||
nn.setCurrentDeviceConfig(device_config)
|
||||
|
||||
# Manipulate environment variables before import tensorflow
|
||||
|
||||
if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
|
||||
os.environ.pop('CUDA_VISIBLE_DEVICES')
|
||||
|
||||
first_run = False
|
||||
if len(device_config.devices) != 0:
|
||||
if sys.platform[0:3] == 'win':
|
||||
# Windows specific env vars
|
||||
if all( [ x.name == device_config.devices[0].name for x in device_config.devices ] ):
|
||||
devices_str = "_" + device_config.devices[0].name.replace(' ','_')
|
||||
else:
|
||||
|
|
@ -123,18 +74,25 @@ class nn():
|
|||
os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2'
|
||||
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # tf log errors only
|
||||
|
||||
import warnings
|
||||
warnings.simplefilter(action='ignore', category=FutureWarning)
|
||||
|
||||
if first_run:
|
||||
io.log_info("Caching GPU kernels...")
|
||||
|
||||
import tensorflow as tf
|
||||
nn.tf = tf
|
||||
|
||||
import logging
|
||||
# Disable tensorflow warnings
|
||||
logging.getLogger('tensorflow').setLevel(logging.ERROR)
|
||||
|
||||
nn.tf = tf
|
||||
|
||||
# Initialize framework
|
||||
import core.leras.ops
|
||||
import core.leras.layers
|
||||
import core.leras.initializers
|
||||
import core.leras.optimizers
|
||||
import core.leras.models
|
||||
import core.leras.archis
|
||||
|
||||
# Configure tensorflow session-config
|
||||
if len(device_config.devices) == 0:
|
||||
nn.tf_default_device = "/CPU:0"
|
||||
config = tf.ConfigProto(device_count={'GPU': 0})
|
||||
|
|
@ -146,20 +104,6 @@ class nn():
|
|||
config.gpu_options.force_gpu_compatible = True
|
||||
config.gpu_options.allow_growth = True
|
||||
nn.tf_sess_config = config
|
||||
|
||||
from .tensor_ops import initialize_tensor_ops
|
||||
from .layers import initialize_layers
|
||||
from .initializers import initialize_initializers
|
||||
from .optimizers import initialize_optimizers
|
||||
from .models import initialize_models
|
||||
from .archis import initialize_archis
|
||||
|
||||
initialize_tensor_ops(nn)
|
||||
initialize_layers(nn)
|
||||
initialize_initializers(nn)
|
||||
initialize_optimizers(nn)
|
||||
initialize_models(nn)
|
||||
initialize_archis(nn)
|
||||
|
||||
if nn.tf_sess is None:
|
||||
nn.tf_sess = tf.Session(config=nn.tf_sess_config)
|
||||
|
|
@ -182,8 +126,7 @@ class nn():
|
|||
"""
|
||||
set default float type for all layers when dtype is None for them
|
||||
"""
|
||||
nn.tf_floatx = tf_dtype
|
||||
nn.np_floatx = tf_dtype.as_numpy_dtype
|
||||
nn.floatx = tf_dtype
|
||||
|
||||
@staticmethod
|
||||
def set_data_format(data_format):
|
||||
|
|
@ -231,7 +174,7 @@ class nn():
|
|||
nn.current_DeviceConfig = device_config
|
||||
|
||||
@staticmethod
|
||||
def tf_reset_session():
|
||||
def reset_session():
|
||||
if nn.tf is not None:
|
||||
if nn.tf_sess is not None:
|
||||
nn.tf.reset_default_graph()
|
||||
|
|
@ -239,14 +182,14 @@ class nn():
|
|||
nn.tf_sess = nn.tf.Session(config=nn.tf_sess_config)
|
||||
|
||||
@staticmethod
|
||||
def tf_close_session():
|
||||
def close_session():
|
||||
if nn.tf_sess is not None:
|
||||
nn.tf.reset_default_graph()
|
||||
nn.tf_sess.close()
|
||||
nn.tf_sess = None
|
||||
|
||||
@staticmethod
|
||||
def tf_get_current_device():
|
||||
def get_current_device():
|
||||
# Undocumented access to last tf.device(...)
|
||||
objs = nn.tf.get_default_graph()._device_function_stack.peek_objs()
|
||||
if len(objs) != 0:
|
||||
|
|
@ -254,7 +197,7 @@ class nn():
|
|||
return nn.tf_default_device
|
||||
|
||||
@staticmethod
|
||||
def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False, return_device_config=False):
|
||||
def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False):
|
||||
devices = Devices.getDevices()
|
||||
if len(devices) == 0:
|
||||
return []
|
||||
|
|
@ -310,12 +253,13 @@ class nn():
|
|||
pass
|
||||
io.log_info ("")
|
||||
|
||||
if return_device_config:
|
||||
return nn.DeviceConfig.GPUIndexes(choosed_idxs)
|
||||
else:
|
||||
return choosed_idxs
|
||||
return choosed_idxs
|
||||
|
||||
class DeviceConfig():
|
||||
@staticmethod
|
||||
def ask_choose_device(*args, **kwargs):
|
||||
return nn.DeviceConfig.GPUIndexes( nn.ask_choose_device_idxs(*args,**kwargs) )
|
||||
|
||||
def __init__ (self, devices=None):
|
||||
devices = devices or []
|
||||
|
||||
|
|
|
|||
345
core/leras/ops/__init__.py
Normal file
345
core/leras/ops/__init__.py
Normal file
|
|
@ -0,0 +1,345 @@
|
|||
import numpy as np
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
from tensorflow.python.ops import array_ops, random_ops, math_ops, sparse_ops, gradients
|
||||
from tensorflow.python.framework import sparse_tensor
|
||||
|
||||
def tf_get_value(tensor):
|
||||
return nn.tf_sess.run (tensor)
|
||||
nn.tf_get_value = tf_get_value
|
||||
|
||||
|
||||
def batch_set_value(tuples):
|
||||
if len(tuples) != 0:
|
||||
with nn.tf.device('/CPU:0'):
|
||||
assign_ops = []
|
||||
feed_dict = {}
|
||||
|
||||
for x, value in tuples:
|
||||
if isinstance(value, nn.tf.Operation) or \
|
||||
isinstance(value, nn.tf.Variable):
|
||||
assign_ops.append(value)
|
||||
else:
|
||||
value = np.asarray(value, dtype=x.dtype.as_numpy_dtype)
|
||||
assign_placeholder = nn.tf.placeholder( x.dtype.base_dtype, shape=[None]*value.ndim )
|
||||
assign_op = nn.tf.assign (x, assign_placeholder )
|
||||
assign_ops.append(assign_op)
|
||||
feed_dict[assign_placeholder] = value
|
||||
|
||||
nn.tf_sess.run(assign_ops, feed_dict=feed_dict)
|
||||
nn.batch_set_value = batch_set_value
|
||||
|
||||
def init_weights(weights):
|
||||
ops = []
|
||||
|
||||
ca_tuples_w = []
|
||||
ca_tuples = []
|
||||
for w in weights:
|
||||
initializer = w.initializer
|
||||
for input in initializer.inputs:
|
||||
if "_cai_" in input.name:
|
||||
ca_tuples_w.append (w)
|
||||
ca_tuples.append ( (w.shape.as_list(), w.dtype.as_numpy_dtype) )
|
||||
break
|
||||
else:
|
||||
ops.append (initializer)
|
||||
|
||||
if len(ops) != 0:
|
||||
nn.tf_sess.run (ops)
|
||||
|
||||
if len(ca_tuples) != 0:
|
||||
nn.batch_set_value( [*zip(ca_tuples_w, nn.initializers.ca.generate_batch (ca_tuples))] )
|
||||
nn.init_weights = init_weights
|
||||
|
||||
def tf_gradients ( loss, vars ):
|
||||
grads = gradients.gradients(loss, vars, colocate_gradients_with_ops=True )
|
||||
gv = [*zip(grads,vars)]
|
||||
for g,v in gv:
|
||||
if g is None:
|
||||
raise Exception(f"No gradient for variable {v.name}")
|
||||
return gv
|
||||
nn.gradients = tf_gradients
|
||||
|
||||
def average_gv_list(grad_var_list, tf_device_string=None):
|
||||
if len(grad_var_list) == 1:
|
||||
return grad_var_list[0]
|
||||
|
||||
e = tf.device(tf_device_string) if tf_device_string is not None else None
|
||||
if e is not None: e.__enter__()
|
||||
result = []
|
||||
for i, (gv) in enumerate(grad_var_list):
|
||||
for j,(g,v) in enumerate(gv):
|
||||
g = tf.expand_dims(g, 0)
|
||||
if i == 0:
|
||||
result += [ [[g], v] ]
|
||||
else:
|
||||
result[j][0] += [g]
|
||||
|
||||
for i,(gs,v) in enumerate(result):
|
||||
result[i] = ( tf.reduce_mean( tf.concat (gs, 0), 0 ), v )
|
||||
if e is not None: e.__exit__(None,None,None)
|
||||
return result
|
||||
nn.average_gv_list = average_gv_list
|
||||
|
||||
def average_tensor_list(tensors_list, tf_device_string=None):
|
||||
if len(tensors_list) == 1:
|
||||
return tensors_list[0]
|
||||
|
||||
e = tf.device(tf_device_string) if tf_device_string is not None else None
|
||||
if e is not None: e.__enter__()
|
||||
result = tf.reduce_mean(tf.concat ([tf.expand_dims(t, 0) for t in tensors_list], 0), 0)
|
||||
if e is not None: e.__exit__(None,None,None)
|
||||
return result
|
||||
nn.average_tensor_list = average_tensor_list
|
||||
|
||||
def concat (tensors_list, axis):
|
||||
"""
|
||||
Better version.
|
||||
"""
|
||||
if len(tensors_list) == 1:
|
||||
return tensors_list[0]
|
||||
return tf.concat(tensors_list, axis)
|
||||
nn.concat = concat
|
||||
|
||||
def gelu(x):
|
||||
cdf = 0.5 * (1.0 + tf.nn.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
|
||||
return x * cdf
|
||||
nn.gelu = gelu
|
||||
|
||||
def upsample2d(x, size=2):
|
||||
if nn.data_format == "NCHW":
|
||||
b,c,h,w = x.shape.as_list()
|
||||
x = tf.reshape (x, (-1,c,h,1,w,1) )
|
||||
x = tf.tile(x, (1,1,1,size,1,size) )
|
||||
x = tf.reshape (x, (-1,c,h*size,w*size) )
|
||||
return x
|
||||
else:
|
||||
return tf.image.resize_nearest_neighbor(x, (x.shape[1]*size, x.shape[2]*size) )
|
||||
nn.upsample2d = upsample2d
|
||||
|
||||
def resize2d_bilinear(x, size=2):
|
||||
h = x.shape[nn.conv2d_spatial_axes[0]].value
|
||||
w = x.shape[nn.conv2d_spatial_axes[1]].value
|
||||
|
||||
if nn.data_format == "NCHW":
|
||||
x = tf.transpose(x, (0,2,3,1))
|
||||
|
||||
if size > 0:
|
||||
new_size = (h*size,w*size)
|
||||
else:
|
||||
new_size = (h//-size,w//-size)
|
||||
|
||||
x = tf.image.resize(x, new_size, method=tf.image.ResizeMethod.BILINEAR)
|
||||
|
||||
if nn.data_format == "NCHW":
|
||||
x = tf.transpose(x, (0,3,1,2))
|
||||
|
||||
return x
|
||||
nn.resize2d_bilinear = resize2d_bilinear
|
||||
|
||||
|
||||
|
||||
def flatten(x):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
x = tf.transpose(x, (0,3,1,2) )
|
||||
return tf.reshape (x, (-1, np.prod(x.shape[1:])) )
|
||||
|
||||
nn.flatten = flatten
|
||||
|
||||
def max_pool(x, kernel_size=2, strides=2):
|
||||
if nn.data_format == "NHWC":
|
||||
return tf.nn.max_pool(x, [1,kernel_size,kernel_size,1], [1,strides,strides,1], 'SAME', data_format=nn.data_format)
|
||||
else:
|
||||
return tf.nn.max_pool(x, [1,1,kernel_size,kernel_size], [1,1,strides,strides], 'SAME', data_format=nn.data_format)
|
||||
|
||||
nn.max_pool = max_pool
|
||||
|
||||
def reshape_4D(x, w,h,c):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
x = tf.reshape (x, (-1,c,h,w))
|
||||
x = tf.transpose(x, (0,2,3,1) )
|
||||
return x
|
||||
else:
|
||||
return tf.reshape (x, (-1,c,h,w))
|
||||
nn.reshape_4D = reshape_4D
|
||||
|
||||
def random_binomial(shape, p=0.0, dtype=None, seed=None):
|
||||
if dtype is None:
|
||||
dtype=tf.float32
|
||||
|
||||
if seed is None:
|
||||
seed = np.random.randint(10e6)
|
||||
return array_ops.where(
|
||||
random_ops.random_uniform(shape, dtype=tf.float16, seed=seed) < p,
|
||||
array_ops.ones(shape, dtype=dtype), array_ops.zeros(shape, dtype=dtype))
|
||||
nn.random_binomial = random_binomial
|
||||
|
||||
def gaussian_blur(input, radius=2.0):
|
||||
def gaussian(x, mu, sigma):
|
||||
return np.exp(-(float(x) - float(mu)) ** 2 / (2 * sigma ** 2))
|
||||
|
||||
def make_kernel(sigma):
|
||||
kernel_size = max(3, int(2 * 2 * sigma + 1))
|
||||
mean = np.floor(0.5 * kernel_size)
|
||||
kernel_1d = np.array([gaussian(x, mean, sigma) for x in range(kernel_size)])
|
||||
np_kernel = np.outer(kernel_1d, kernel_1d).astype(np.float32)
|
||||
kernel = np_kernel / np.sum(np_kernel)
|
||||
return kernel, kernel_size
|
||||
|
||||
gauss_kernel, kernel_size = make_kernel(radius)
|
||||
padding = kernel_size//2
|
||||
if padding != 0:
|
||||
if nn.data_format == "NHWC":
|
||||
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
|
||||
else:
|
||||
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
|
||||
else:
|
||||
padding = None
|
||||
gauss_kernel = gauss_kernel[:,:,None,None]
|
||||
|
||||
x = input
|
||||
k = tf.tile (gauss_kernel, (1,1,x.shape[nn.conv2d_ch_axis],1) )
|
||||
x = tf.pad(x, padding )
|
||||
x = tf.nn.depthwise_conv2d(x, k, strides=[1,1,1,1], padding='VALID', data_format=nn.data_format)
|
||||
return x
|
||||
nn.gaussian_blur = gaussian_blur
|
||||
|
||||
def style_loss(target, style, gaussian_blur_radius=0.0, loss_weight=1.0, step_size=1):
|
||||
def sd(content, style, loss_weight):
|
||||
content_nc = content.shape[ nn.conv2d_ch_axis ]
|
||||
style_nc = style.shape[nn.conv2d_ch_axis]
|
||||
if content_nc != style_nc:
|
||||
raise Exception("style_loss() content_nc != style_nc")
|
||||
c_mean, c_var = tf.nn.moments(content, axes=nn.conv2d_spatial_axes, keep_dims=True)
|
||||
s_mean, s_var = tf.nn.moments(style, axes=nn.conv2d_spatial_axes, keep_dims=True)
|
||||
c_std, s_std = tf.sqrt(c_var + 1e-5), tf.sqrt(s_var + 1e-5)
|
||||
mean_loss = tf.reduce_sum(tf.square(c_mean-s_mean), axis=[1,2,3])
|
||||
std_loss = tf.reduce_sum(tf.square(c_std-s_std), axis=[1,2,3])
|
||||
return (mean_loss + std_loss) * ( loss_weight / content_nc.value )
|
||||
|
||||
if gaussian_blur_radius > 0.0:
|
||||
target = gaussian_blur(target, gaussian_blur_radius)
|
||||
style = gaussian_blur(style, gaussian_blur_radius)
|
||||
|
||||
return sd( target, style, loss_weight=loss_weight )
|
||||
|
||||
nn.style_loss = style_loss
|
||||
|
||||
def dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03):
|
||||
if img1.dtype != img2.dtype:
|
||||
raise ValueError("img1.dtype != img2.dtype")
|
||||
|
||||
not_float32 = img1.dtype != tf.float32
|
||||
|
||||
if not_float32:
|
||||
img_dtype = img1.dtype
|
||||
img1 = tf.cast(img1, tf.float32)
|
||||
img2 = tf.cast(img2, tf.float32)
|
||||
|
||||
kernel = np.arange(0, filter_size, dtype=np.float32)
|
||||
kernel -= (filter_size - 1 ) / 2.0
|
||||
kernel = kernel**2
|
||||
kernel *= ( -0.5 / (filter_sigma**2) )
|
||||
kernel = np.reshape (kernel, (1,-1)) + np.reshape(kernel, (-1,1) )
|
||||
kernel = tf.constant ( np.reshape (kernel, (1,-1)), dtype=tf.float32 )
|
||||
kernel = tf.nn.softmax(kernel)
|
||||
kernel = tf.reshape (kernel, (filter_size, filter_size, 1, 1))
|
||||
kernel = tf.tile (kernel, (1,1, img1.shape[ nn.conv2d_ch_axis ] ,1))
|
||||
|
||||
def reducer(x):
|
||||
return tf.nn.depthwise_conv2d(x, kernel, strides=[1,1,1,1], padding='VALID', data_format=nn.data_format)
|
||||
|
||||
c1 = (k1 * max_val) ** 2
|
||||
c2 = (k2 * max_val) ** 2
|
||||
|
||||
mean0 = reducer(img1)
|
||||
mean1 = reducer(img2)
|
||||
num0 = mean0 * mean1 * 2.0
|
||||
den0 = tf.square(mean0) + tf.square(mean1)
|
||||
luminance = (num0 + c1) / (den0 + c1)
|
||||
|
||||
num1 = reducer(img1 * img2) * 2.0
|
||||
den1 = reducer(tf.square(img1) + tf.square(img2))
|
||||
c2 *= 1.0 #compensation factor
|
||||
cs = (num1 - num0 + c2) / (den1 - den0 + c2)
|
||||
|
||||
ssim_val = tf.reduce_mean(luminance * cs, axis=nn.conv2d_spatial_axes )
|
||||
dssim = (1.0 - ssim_val ) / 2.0
|
||||
|
||||
if not_float32:
|
||||
dssim = tf.cast(dssim, img_dtype)
|
||||
return dssim
|
||||
|
||||
nn.dssim = dssim
|
||||
|
||||
def space_to_depth(x, size):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
b,h,w,c = x.shape.as_list()
|
||||
oh, ow = h // size, w // size
|
||||
x = tf.reshape(x, (-1, size, oh, size, ow, c))
|
||||
x = tf.transpose(x, (0, 2, 4, 1, 3, 5))
|
||||
x = tf.reshape(x, (-1, oh, ow, size* size* c ))
|
||||
return x
|
||||
else:
|
||||
return tf.space_to_depth(x, size, data_format=nn.data_format)
|
||||
nn.space_to_depth = space_to_depth
|
||||
|
||||
def depth_to_space(x, size):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
|
||||
b,h,w,c = x.shape.as_list()
|
||||
oh, ow = h * size, w * size
|
||||
oc = c // (size * size)
|
||||
|
||||
x = tf.reshape(x, (-1, h, w, size, size, oc, ) )
|
||||
x = tf.transpose(x, (0, 1, 3, 2, 4, 5))
|
||||
x = tf.reshape(x, (-1, oh, ow, oc, ))
|
||||
return x
|
||||
else:
|
||||
return tf.depth_to_space(x, size, data_format=nn.data_format)
|
||||
nn.depth_to_space = depth_to_space
|
||||
|
||||
def rgb_to_lab(srgb):
|
||||
srgb_pixels = tf.reshape(srgb, [-1, 3])
|
||||
linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)
|
||||
exponential_mask = tf.cast(srgb_pixels > 0.04045, dtype=tf.float32)
|
||||
rgb_pixels = (srgb_pixels / 12.92 * linear_mask) + (((srgb_pixels + 0.055) / 1.055) ** 2.4) * exponential_mask
|
||||
rgb_to_xyz = tf.constant([
|
||||
# X Y Z
|
||||
[0.412453, 0.212671, 0.019334], # R
|
||||
[0.357580, 0.715160, 0.119193], # G
|
||||
[0.180423, 0.072169, 0.950227], # B
|
||||
])
|
||||
xyz_pixels = tf.matmul(rgb_pixels, rgb_to_xyz)
|
||||
|
||||
xyz_normalized_pixels = tf.multiply(xyz_pixels, [1/0.950456, 1.0, 1/1.088754])
|
||||
|
||||
epsilon = 6/29
|
||||
linear_mask = tf.cast(xyz_normalized_pixels <= (epsilon**3), dtype=tf.float32)
|
||||
exponential_mask = tf.cast(xyz_normalized_pixels > (epsilon**3), dtype=tf.float32)
|
||||
fxfyfz_pixels = (xyz_normalized_pixels / (3 * epsilon**2) + 4/29) * linear_mask + (xyz_normalized_pixels ** (1/3)) * exponential_mask
|
||||
|
||||
fxfyfz_to_lab = tf.constant([
|
||||
# l a b
|
||||
[ 0.0, 500.0, 0.0], # fx
|
||||
[116.0, -500.0, 200.0], # fy
|
||||
[ 0.0, 0.0, -200.0], # fz
|
||||
])
|
||||
lab_pixels = tf.matmul(fxfyfz_pixels, fxfyfz_to_lab) + tf.constant([-16.0, 0.0, 0.0])
|
||||
return tf.reshape(lab_pixels, tf.shape(srgb))
|
||||
nn.rgb_to_lab = rgb_to_lab
|
||||
|
||||
"""
|
||||
def tf_suppress_lower_mean(t, eps=0.00001):
|
||||
if t.shape.ndims != 1:
|
||||
raise ValueError("tf_suppress_lower_mean: t rank must be 1")
|
||||
t_mean_eps = tf.reduce_mean(t) - eps
|
||||
q = tf.clip_by_value(t, t_mean_eps, tf.reduce_max(t) )
|
||||
q = tf.clip_by_value(q-t_mean_eps, 0, eps)
|
||||
q = q * (t/eps)
|
||||
return q
|
||||
"""
|
||||
|
|
@ -1,110 +0,0 @@
|
|||
import copy
|
||||
|
||||
def initialize_optimizers(nn):
|
||||
tf = nn.tf
|
||||
from tensorflow.python.ops import state_ops, control_flow_ops
|
||||
|
||||
class TFBaseOptimizer(nn.Saveable):
|
||||
def __init__(self, name=None):
|
||||
super().__init__(name=name)
|
||||
|
||||
def tf_clip_norm(self, g, c, n):
|
||||
"""Clip the gradient `g` if the L2 norm `n` exceeds `c`.
|
||||
# Arguments
|
||||
g: Tensor, the gradient tensor
|
||||
c: float >= 0. Gradients will be clipped
|
||||
when their L2 norm exceeds this value.
|
||||
n: Tensor, actual norm of `g`.
|
||||
# Returns
|
||||
Tensor, the gradient clipped if required.
|
||||
"""
|
||||
if c <= 0: # if clipnorm == 0 no need to add ops to the graph
|
||||
return g
|
||||
|
||||
condition = n >= c
|
||||
then_expression = tf.scalar_mul(c / n, g)
|
||||
else_expression = g
|
||||
|
||||
# saving the shape to avoid converting sparse tensor to dense
|
||||
if isinstance(then_expression, tf.Tensor):
|
||||
g_shape = copy.copy(then_expression.get_shape())
|
||||
elif isinstance(then_expression, tf.IndexedSlices):
|
||||
g_shape = copy.copy(then_expression.dense_shape)
|
||||
if condition.dtype != tf.bool:
|
||||
condition = tf.cast(condition, 'bool')
|
||||
g = tf.cond(condition,
|
||||
lambda: then_expression,
|
||||
lambda: else_expression)
|
||||
if isinstance(then_expression, tf.Tensor):
|
||||
g.set_shape(g_shape)
|
||||
elif isinstance(then_expression, tf.IndexedSlices):
|
||||
g._dense_shape = g_shape
|
||||
|
||||
return g
|
||||
nn.TFBaseOptimizer = TFBaseOptimizer
|
||||
|
||||
class TFRMSpropOptimizer(TFBaseOptimizer):
|
||||
def __init__(self, lr=0.001, rho=0.9, lr_dropout=1.0, epsilon=1e-7, clipnorm=0.0, name=None):
|
||||
super().__init__(name=name)
|
||||
|
||||
if name is None:
|
||||
raise ValueError('name must be defined.')
|
||||
|
||||
self.lr_dropout = lr_dropout
|
||||
self.clipnorm = clipnorm
|
||||
|
||||
with tf.device('/CPU:0') :
|
||||
with tf.variable_scope(self.name):
|
||||
self.lr = tf.Variable (lr, name="lr")
|
||||
self.rho = tf.Variable (rho, name="rho")
|
||||
self.epsilon = tf.Variable (epsilon, name="epsilon")
|
||||
self.iterations = tf.Variable(0, dtype=tf.int64, name='iters')
|
||||
|
||||
self.accumulators_dict = {}
|
||||
self.lr_rnds_dict = {}
|
||||
|
||||
def get_weights(self):
|
||||
return [self.lr, self.rho, self.epsilon, self.iterations] + list(self.accumulators_dict.values())
|
||||
|
||||
def initialize_variables(self, trainable_weights, vars_on_cpu=True):
|
||||
# Initialize here all trainable variables used in training
|
||||
e = tf.device('/CPU:0') if vars_on_cpu else None
|
||||
if e: e.__enter__()
|
||||
with tf.variable_scope(self.name):
|
||||
accumulators = { v.name : tf.get_variable ( f'acc_{v.name}'.replace(':','_'), v.shape, dtype=v.dtype, initializer=tf.initializers.constant(0.0), trainable=False) for v in trainable_weights }
|
||||
self.accumulators_dict.update ( accumulators)
|
||||
|
||||
if self.lr_dropout != 1.0:
|
||||
lr_rnds = [ nn.tf_random_binomial( v.shape, p=self.lr_dropout, dtype=v.dtype) for v in trainable_weights ]
|
||||
self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } )
|
||||
if e: e.__exit__(None, None, None)
|
||||
|
||||
def get_update_op(self, grads_vars):
|
||||
updates = []
|
||||
|
||||
if self.clipnorm > 0.0:
|
||||
norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars]))
|
||||
updates += [ state_ops.assign_add( self.iterations, 1) ]
|
||||
for i, (g,v) in enumerate(grads_vars):
|
||||
if self.clipnorm > 0.0:
|
||||
g = self.tf_clip_norm(g, self.clipnorm, norm)
|
||||
|
||||
a = self.accumulators_dict[ v.name ]
|
||||
|
||||
rho = tf.cast(self.rho, a.dtype)
|
||||
new_a = rho * a + (1. - rho) * tf.square(g)
|
||||
|
||||
lr = tf.cast(self.lr, a.dtype)
|
||||
epsilon = tf.cast(self.epsilon, a.dtype)
|
||||
|
||||
v_diff = - lr * g / (tf.sqrt(new_a) + epsilon)
|
||||
if self.lr_dropout != 1.0:
|
||||
lr_rnd = self.lr_rnds_dict[v.name]
|
||||
v_diff *= lr_rnd
|
||||
new_v = v + v_diff
|
||||
|
||||
updates.append (state_ops.assign(a, new_a))
|
||||
updates.append (state_ops.assign(v, new_v))
|
||||
|
||||
return control_flow_ops.group ( *updates, name=self.name+'_updates')
|
||||
nn.TFRMSpropOptimizer = TFRMSpropOptimizer
|
||||
42
core/leras/optimizers/OptimizerBase.py
Normal file
42
core/leras/optimizers/OptimizerBase.py
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
import copy
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class OptimizerBase(nn.Saveable):
|
||||
def __init__(self, name=None):
|
||||
super().__init__(name=name)
|
||||
|
||||
def tf_clip_norm(self, g, c, n):
|
||||
"""Clip the gradient `g` if the L2 norm `n` exceeds `c`.
|
||||
# Arguments
|
||||
g: Tensor, the gradient tensor
|
||||
c: float >= 0. Gradients will be clipped
|
||||
when their L2 norm exceeds this value.
|
||||
n: Tensor, actual norm of `g`.
|
||||
# Returns
|
||||
Tensor, the gradient clipped if required.
|
||||
"""
|
||||
if c <= 0: # if clipnorm == 0 no need to add ops to the graph
|
||||
return g
|
||||
|
||||
condition = n >= c
|
||||
then_expression = tf.scalar_mul(c / n, g)
|
||||
else_expression = g
|
||||
|
||||
# saving the shape to avoid converting sparse tensor to dense
|
||||
if isinstance(then_expression, tf.Tensor):
|
||||
g_shape = copy.copy(then_expression.get_shape())
|
||||
elif isinstance(then_expression, tf.IndexedSlices):
|
||||
g_shape = copy.copy(then_expression.dense_shape)
|
||||
if condition.dtype != tf.bool:
|
||||
condition = tf.cast(condition, 'bool')
|
||||
g = tf.cond(condition,
|
||||
lambda: then_expression,
|
||||
lambda: else_expression)
|
||||
if isinstance(then_expression, tf.Tensor):
|
||||
g.set_shape(g_shape)
|
||||
elif isinstance(then_expression, tf.IndexedSlices):
|
||||
g._dense_shape = g_shape
|
||||
|
||||
return g
|
||||
nn.OptimizerBase = OptimizerBase
|
||||
69
core/leras/optimizers/RMSprop.py
Normal file
69
core/leras/optimizers/RMSprop.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
from tensorflow.python.ops import control_flow_ops, state_ops
|
||||
from core.leras import nn
|
||||
tf = nn.tf
|
||||
|
||||
class RMSprop(nn.OptimizerBase):
|
||||
def __init__(self, lr=0.001, rho=0.9, lr_dropout=1.0, epsilon=1e-7, clipnorm=0.0, name=None):
|
||||
super().__init__(name=name)
|
||||
|
||||
if name is None:
|
||||
raise ValueError('name must be defined.')
|
||||
|
||||
self.lr_dropout = lr_dropout
|
||||
self.clipnorm = clipnorm
|
||||
|
||||
with tf.device('/CPU:0') :
|
||||
with tf.variable_scope(self.name):
|
||||
self.lr = tf.Variable (lr, name="lr")
|
||||
self.rho = tf.Variable (rho, name="rho")
|
||||
self.epsilon = tf.Variable (epsilon, name="epsilon")
|
||||
self.iterations = tf.Variable(0, dtype=tf.int64, name='iters')
|
||||
|
||||
self.accumulators_dict = {}
|
||||
self.lr_rnds_dict = {}
|
||||
|
||||
def get_weights(self):
|
||||
return [self.lr, self.rho, self.epsilon, self.iterations] + list(self.accumulators_dict.values())
|
||||
|
||||
def initialize_variables(self, trainable_weights, vars_on_cpu=True):
|
||||
# Initialize here all trainable variables used in training
|
||||
e = tf.device('/CPU:0') if vars_on_cpu else None
|
||||
if e: e.__enter__()
|
||||
with tf.variable_scope(self.name):
|
||||
accumulators = { v.name : tf.get_variable ( f'acc_{v.name}'.replace(':','_'), v.shape, dtype=v.dtype, initializer=tf.initializers.constant(0.0), trainable=False) for v in trainable_weights }
|
||||
self.accumulators_dict.update ( accumulators)
|
||||
|
||||
if self.lr_dropout != 1.0:
|
||||
lr_rnds = [ nn.random_binomial( v.shape, p=self.lr_dropout, dtype=v.dtype) for v in trainable_weights ]
|
||||
self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } )
|
||||
if e: e.__exit__(None, None, None)
|
||||
|
||||
def get_update_op(self, grads_vars):
|
||||
updates = []
|
||||
|
||||
if self.clipnorm > 0.0:
|
||||
norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars]))
|
||||
updates += [ state_ops.assign_add( self.iterations, 1) ]
|
||||
for i, (g,v) in enumerate(grads_vars):
|
||||
if self.clipnorm > 0.0:
|
||||
g = self.tf_clip_norm(g, self.clipnorm, norm)
|
||||
|
||||
a = self.accumulators_dict[ v.name ]
|
||||
|
||||
rho = tf.cast(self.rho, a.dtype)
|
||||
new_a = rho * a + (1. - rho) * tf.square(g)
|
||||
|
||||
lr = tf.cast(self.lr, a.dtype)
|
||||
epsilon = tf.cast(self.epsilon, a.dtype)
|
||||
|
||||
v_diff = - lr * g / (tf.sqrt(new_a) + epsilon)
|
||||
if self.lr_dropout != 1.0:
|
||||
lr_rnd = self.lr_rnds_dict[v.name]
|
||||
v_diff *= lr_rnd
|
||||
new_v = v + v_diff
|
||||
|
||||
updates.append (state_ops.assign(a, new_a))
|
||||
updates.append (state_ops.assign(v, new_v))
|
||||
|
||||
return control_flow_ops.group ( *updates, name=self.name+'_updates')
|
||||
nn.RMSprop = RMSprop
|
||||
2
core/leras/optimizers/__init__.py
Normal file
2
core/leras/optimizers/__init__.py
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
from .OptimizerBase import *
|
||||
from .RMSprop import *
|
||||
|
|
@ -1,374 +0,0 @@
|
|||
import numpy as np
|
||||
|
||||
def initialize_tensor_ops(nn):
|
||||
tf = nn.tf
|
||||
from tensorflow.python.ops import array_ops, random_ops, math_ops, sparse_ops, gradients
|
||||
from tensorflow.python.framework import sparse_tensor
|
||||
|
||||
def tf_get_value(tensor):
|
||||
return nn.tf_sess.run (tensor)
|
||||
nn.tf_get_value = tf_get_value
|
||||
|
||||
|
||||
def tf_batch_set_value(tuples):
|
||||
if len(tuples) != 0:
|
||||
with nn.tf.device('/CPU:0'):
|
||||
assign_ops = []
|
||||
feed_dict = {}
|
||||
|
||||
for x, value in tuples:
|
||||
if isinstance(value, nn.tf.Operation):
|
||||
assign_ops.append(value)
|
||||
else:
|
||||
value = np.asarray(value, dtype=x.dtype.as_numpy_dtype)
|
||||
assign_placeholder = nn.tf.placeholder( x.dtype.base_dtype, shape=[None]*value.ndim )
|
||||
assign_op = nn.tf.assign (x, assign_placeholder )
|
||||
assign_ops.append(assign_op)
|
||||
feed_dict[assign_placeholder] = value
|
||||
|
||||
nn.tf_sess.run(assign_ops, feed_dict=feed_dict)
|
||||
nn.tf_batch_set_value = tf_batch_set_value
|
||||
|
||||
def tf_init_weights(weights):
|
||||
ops = []
|
||||
|
||||
ca_tuples_w = []
|
||||
ca_tuples = []
|
||||
for w in weights:
|
||||
initializer = w.initializer
|
||||
for input in initializer.inputs:
|
||||
if "_cai_" in input.name:
|
||||
ca_tuples_w.append (w)
|
||||
ca_tuples.append ( (w.shape.as_list(), w.dtype.as_numpy_dtype) )
|
||||
break
|
||||
else:
|
||||
ops.append (initializer)
|
||||
|
||||
if len(ops) != 0:
|
||||
nn.tf_sess.run (ops)
|
||||
|
||||
if len(ca_tuples) != 0:
|
||||
nn.tf_batch_set_value( [*zip(ca_tuples_w, nn.initializers.ca.generate_batch (ca_tuples))] )
|
||||
nn.tf_init_weights = tf_init_weights
|
||||
|
||||
def tf_gradients ( loss, vars ):
|
||||
grads = gradients.gradients(loss, vars, colocate_gradients_with_ops=True )
|
||||
gv = [*zip(grads,vars)]
|
||||
for g,v in gv:
|
||||
if g is None:
|
||||
raise Exception(f"No gradient for variable {v.name}")
|
||||
return gv
|
||||
nn.tf_gradients = tf_gradients
|
||||
|
||||
def tf_average_gv_list(grad_var_list, tf_device_string=None):
|
||||
if len(grad_var_list) == 1:
|
||||
return grad_var_list[0]
|
||||
|
||||
e = tf.device(tf_device_string) if tf_device_string is not None else None
|
||||
if e is not None: e.__enter__()
|
||||
result = []
|
||||
for i, (gv) in enumerate(grad_var_list):
|
||||
for j,(g,v) in enumerate(gv):
|
||||
g = tf.expand_dims(g, 0)
|
||||
if i == 0:
|
||||
result += [ [[g], v] ]
|
||||
else:
|
||||
result[j][0] += [g]
|
||||
|
||||
for i,(gs,v) in enumerate(result):
|
||||
result[i] = ( tf.reduce_mean( tf.concat (gs, 0), 0 ), v )
|
||||
if e is not None: e.__exit__(None,None,None)
|
||||
return result
|
||||
nn.tf_average_gv_list = tf_average_gv_list
|
||||
|
||||
def tf_average_tensor_list(tensors_list, tf_device_string=None):
|
||||
if len(tensors_list) == 1:
|
||||
return tensors_list[0]
|
||||
|
||||
e = tf.device(tf_device_string) if tf_device_string is not None else None
|
||||
if e is not None: e.__enter__()
|
||||
result = tf.reduce_mean(tf.concat ([tf.expand_dims(t, 0) for t in tensors_list], 0), 0)
|
||||
if e is not None: e.__exit__(None,None,None)
|
||||
return result
|
||||
nn.tf_average_tensor_list = tf_average_tensor_list
|
||||
|
||||
def tf_concat (tensors_list, axis):
|
||||
"""
|
||||
Better version.
|
||||
"""
|
||||
if len(tensors_list) == 1:
|
||||
return tensors_list[0]
|
||||
return tf.concat(tensors_list, axis)
|
||||
nn.tf_concat = tf_concat
|
||||
|
||||
def tf_gelu(x):
|
||||
cdf = 0.5 * (1.0 + tf.nn.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
|
||||
return x * cdf
|
||||
nn.tf_gelu = tf_gelu
|
||||
|
||||
def tf_upsample2d(x, size=2):
|
||||
if nn.data_format == "NCHW":
|
||||
b,c,h,w = x.shape.as_list()
|
||||
x = tf.reshape (x, (-1,c,h,1,w,1) )
|
||||
x = tf.tile(x, (1,1,1,size,1,size) )
|
||||
x = tf.reshape (x, (-1,c,h*size,w*size) )
|
||||
return x
|
||||
else:
|
||||
return tf.image.resize_nearest_neighbor(x, (x.shape[1]*size, x.shape[2]*size) )
|
||||
nn.tf_upsample2d = tf_upsample2d
|
||||
|
||||
def tf_resize2d_bilinear(x, size=2):
|
||||
h = x.shape[nn.conv2d_spatial_axes[0]].value
|
||||
w = x.shape[nn.conv2d_spatial_axes[1]].value
|
||||
|
||||
if nn.data_format == "NCHW":
|
||||
x = tf.transpose(x, (0,2,3,1))
|
||||
|
||||
if size > 0:
|
||||
new_size = (h*size,w*size)
|
||||
else:
|
||||
new_size = (h//-size,w//-size)
|
||||
|
||||
x = tf.image.resize(x, new_size, method=tf.image.ResizeMethod.BILINEAR)
|
||||
|
||||
if nn.data_format == "NCHW":
|
||||
x = tf.transpose(x, (0,3,1,2))
|
||||
|
||||
return x
|
||||
nn.tf_resize2d_bilinear = tf_resize2d_bilinear
|
||||
|
||||
|
||||
|
||||
def tf_flatten(x):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
x = tf.transpose(x, (0,3,1,2) )
|
||||
return tf.reshape (x, (-1, np.prod(x.shape[1:])) )
|
||||
|
||||
nn.tf_flatten = tf_flatten
|
||||
|
||||
def tf_max_pool(x, kernel_size, strides):
|
||||
if nn.data_format == "NHWC":
|
||||
return tf.nn.max_pool(x, [1,kernel_size,kernel_size,1], [1,strides,strides,1], "VALID", data_format=nn.data_format)
|
||||
else:
|
||||
return tf.nn.max_pool(x, [1,1,kernel_size,kernel_size], [1,1,strides,strides], "VALID", data_format=nn.data_format)
|
||||
|
||||
nn.tf_max_pool = tf_max_pool
|
||||
|
||||
def tf_reshape_4D(x, w,h,c):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
x = tf.reshape (x, (-1,c,h,w))
|
||||
x = tf.transpose(x, (0,2,3,1) )
|
||||
return x
|
||||
else:
|
||||
return tf.reshape (x, (-1,c,h,w))
|
||||
nn.tf_reshape_4D = tf_reshape_4D
|
||||
|
||||
def tf_random_binomial(shape, p=0.0, dtype=None, seed=None):
|
||||
if dtype is None:
|
||||
dtype=tf.float32
|
||||
|
||||
if seed is None:
|
||||
seed = np.random.randint(10e6)
|
||||
return array_ops.where(
|
||||
random_ops.random_uniform(shape, dtype=tf.float16, seed=seed) < p,
|
||||
array_ops.ones(shape, dtype=dtype), array_ops.zeros(shape, dtype=dtype))
|
||||
nn.tf_random_binomial = tf_random_binomial
|
||||
|
||||
def tf_gaussian_blur(input, radius=2.0):
|
||||
def gaussian(x, mu, sigma):
|
||||
return np.exp(-(float(x) - float(mu)) ** 2 / (2 * sigma ** 2))
|
||||
|
||||
def make_kernel(sigma):
|
||||
kernel_size = max(3, int(2 * 2 * sigma + 1))
|
||||
mean = np.floor(0.5 * kernel_size)
|
||||
kernel_1d = np.array([gaussian(x, mean, sigma) for x in range(kernel_size)])
|
||||
np_kernel = np.outer(kernel_1d, kernel_1d).astype(np.float32)
|
||||
kernel = np_kernel / np.sum(np_kernel)
|
||||
return kernel, kernel_size
|
||||
|
||||
gauss_kernel, kernel_size = make_kernel(radius)
|
||||
padding = kernel_size//2
|
||||
if padding != 0:
|
||||
if nn.data_format == "NHWC":
|
||||
padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
|
||||
else:
|
||||
padding = [ [0,0], [0,0], [padding,padding], [padding,padding] ]
|
||||
else:
|
||||
padding = None
|
||||
gauss_kernel = gauss_kernel[:,:,None,None]
|
||||
|
||||
x = input
|
||||
k = tf.tile (gauss_kernel, (1,1,x.shape[nn.conv2d_ch_axis],1) )
|
||||
x = tf.pad(x, padding )
|
||||
x = tf.nn.depthwise_conv2d(x, k, strides=[1,1,1,1], padding='VALID', data_format=nn.data_format)
|
||||
return x
|
||||
nn.tf_gaussian_blur = tf_gaussian_blur
|
||||
|
||||
def tf_style_loss(target, style, gaussian_blur_radius=0.0, loss_weight=1.0, step_size=1):
|
||||
def sd(content, style, loss_weight):
|
||||
content_nc = content.shape[ nn.conv2d_ch_axis ]
|
||||
style_nc = style.shape[nn.conv2d_ch_axis]
|
||||
if content_nc != style_nc:
|
||||
raise Exception("style_loss() content_nc != style_nc")
|
||||
c_mean, c_var = tf.nn.moments(content, axes=nn.conv2d_spatial_axes, keep_dims=True)
|
||||
s_mean, s_var = tf.nn.moments(style, axes=nn.conv2d_spatial_axes, keep_dims=True)
|
||||
c_std, s_std = tf.sqrt(c_var + 1e-5), tf.sqrt(s_var + 1e-5)
|
||||
mean_loss = tf.reduce_sum(tf.square(c_mean-s_mean), axis=[1,2,3])
|
||||
std_loss = tf.reduce_sum(tf.square(c_std-s_std), axis=[1,2,3])
|
||||
return (mean_loss + std_loss) * ( loss_weight / content_nc.value )
|
||||
|
||||
if gaussian_blur_radius > 0.0:
|
||||
target = tf_gaussian_blur(target, gaussian_blur_radius)
|
||||
style = tf_gaussian_blur(style, gaussian_blur_radius)
|
||||
|
||||
return sd( target, style, loss_weight=loss_weight )
|
||||
|
||||
nn.tf_style_loss = tf_style_loss
|
||||
|
||||
def tf_dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03):
|
||||
if img1.dtype != img2.dtype:
|
||||
raise ValueError("img1.dtype != img2.dtype")
|
||||
|
||||
not_float32 = img1.dtype != tf.float32
|
||||
|
||||
if not_float32:
|
||||
img_dtype = img1.dtype
|
||||
img1 = tf.cast(img1, tf.float32)
|
||||
img2 = tf.cast(img2, tf.float32)
|
||||
|
||||
kernel = np.arange(0, filter_size, dtype=np.float32)
|
||||
kernel -= (filter_size - 1 ) / 2.0
|
||||
kernel = kernel**2
|
||||
kernel *= ( -0.5 / (filter_sigma**2) )
|
||||
kernel = np.reshape (kernel, (1,-1)) + np.reshape(kernel, (-1,1) )
|
||||
kernel = tf.constant ( np.reshape (kernel, (1,-1)), dtype=tf.float32 )
|
||||
kernel = tf.nn.softmax(kernel)
|
||||
kernel = tf.reshape (kernel, (filter_size, filter_size, 1, 1))
|
||||
kernel = tf.tile (kernel, (1,1, img1.shape[ nn.conv2d_ch_axis ] ,1))
|
||||
|
||||
def reducer(x):
|
||||
return tf.nn.depthwise_conv2d(x, kernel, strides=[1,1,1,1], padding='VALID', data_format=nn.data_format)
|
||||
|
||||
c1 = (k1 * max_val) ** 2
|
||||
c2 = (k2 * max_val) ** 2
|
||||
|
||||
mean0 = reducer(img1)
|
||||
mean1 = reducer(img2)
|
||||
num0 = mean0 * mean1 * 2.0
|
||||
den0 = tf.square(mean0) + tf.square(mean1)
|
||||
luminance = (num0 + c1) / (den0 + c1)
|
||||
|
||||
num1 = reducer(img1 * img2) * 2.0
|
||||
den1 = reducer(tf.square(img1) + tf.square(img2))
|
||||
c2 *= 1.0 #compensation factor
|
||||
cs = (num1 - num0 + c2) / (den1 - den0 + c2)
|
||||
|
||||
ssim_val = tf.reduce_mean(luminance * cs, axis=nn.conv2d_spatial_axes )
|
||||
dssim = (1.0 - ssim_val ) / 2.0
|
||||
|
||||
if not_float32:
|
||||
dssim = tf.cast(dssim, img_dtype)
|
||||
return dssim
|
||||
|
||||
nn.tf_dssim = tf_dssim
|
||||
|
||||
def tf_space_to_depth(x, size):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
b,h,w,c = x.shape.as_list()
|
||||
oh, ow = h // size, w // size
|
||||
x = tf.reshape(x, (-1, size, oh, size, ow, c))
|
||||
x = tf.transpose(x, (0, 2, 4, 1, 3, 5))
|
||||
x = tf.reshape(x, (-1, oh, ow, size* size* c ))
|
||||
return x
|
||||
else:
|
||||
return tf.space_to_depth(x, size, data_format=nn.data_format)
|
||||
nn.tf_space_to_depth = tf_space_to_depth
|
||||
|
||||
def tf_depth_to_space(x, size):
|
||||
if nn.data_format == "NHWC":
|
||||
# match NCHW version in order to switch data_format without problems
|
||||
|
||||
b,h,w,c = x.shape.as_list()
|
||||
oh, ow = h * size, w * size
|
||||
oc = c // (size * size)
|
||||
|
||||
x = tf.reshape(x, (-1, h, w, size, size, oc, ) )
|
||||
x = tf.transpose(x, (0, 1, 3, 2, 4, 5))
|
||||
x = tf.reshape(x, (-1, oh, ow, oc, ))
|
||||
return x
|
||||
else:
|
||||
return tf.depth_to_space(x, size, data_format=nn.data_format)
|
||||
nn.tf_depth_to_space = tf_depth_to_space
|
||||
|
||||
def tf_rgb_to_lab(srgb):
|
||||
srgb_pixels = tf.reshape(srgb, [-1, 3])
|
||||
linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)
|
||||
exponential_mask = tf.cast(srgb_pixels > 0.04045, dtype=tf.float32)
|
||||
rgb_pixels = (srgb_pixels / 12.92 * linear_mask) + (((srgb_pixels + 0.055) / 1.055) ** 2.4) * exponential_mask
|
||||
rgb_to_xyz = tf.constant([
|
||||
# X Y Z
|
||||
[0.412453, 0.212671, 0.019334], # R
|
||||
[0.357580, 0.715160, 0.119193], # G
|
||||
[0.180423, 0.072169, 0.950227], # B
|
||||
])
|
||||
xyz_pixels = tf.matmul(rgb_pixels, rgb_to_xyz)
|
||||
|
||||
xyz_normalized_pixels = tf.multiply(xyz_pixels, [1/0.950456, 1.0, 1/1.088754])
|
||||
|
||||
epsilon = 6/29
|
||||
linear_mask = tf.cast(xyz_normalized_pixels <= (epsilon**3), dtype=tf.float32)
|
||||
exponential_mask = tf.cast(xyz_normalized_pixels > (epsilon**3), dtype=tf.float32)
|
||||
fxfyfz_pixels = (xyz_normalized_pixels / (3 * epsilon**2) + 4/29) * linear_mask + (xyz_normalized_pixels ** (1/3)) * exponential_mask
|
||||
|
||||
fxfyfz_to_lab = tf.constant([
|
||||
# l a b
|
||||
[ 0.0, 500.0, 0.0], # fx
|
||||
[116.0, -500.0, 200.0], # fy
|
||||
[ 0.0, 0.0, -200.0], # fz
|
||||
])
|
||||
lab_pixels = tf.matmul(fxfyfz_pixels, fxfyfz_to_lab) + tf.constant([-16.0, 0.0, 0.0])
|
||||
return tf.reshape(lab_pixels, tf.shape(srgb))
|
||||
nn.tf_rgb_to_lab = tf_rgb_to_lab
|
||||
|
||||
def tf_suppress_lower_mean(t, eps=0.00001):
|
||||
if t.shape.ndims != 1:
|
||||
raise ValueError("tf_suppress_lower_mean: t rank must be 1")
|
||||
t_mean_eps = tf.reduce_mean(t) - eps
|
||||
q = tf.clip_by_value(t, t_mean_eps, tf.reduce_max(t) )
|
||||
q = tf.clip_by_value(q-t_mean_eps, 0, eps)
|
||||
q = q * (t/eps)
|
||||
return q
|
||||
"""
|
||||
class GeLU(KL.Layer):
|
||||
Gaussian Error Linear Unit.
|
||||
A smoother version of ReLU generally used
|
||||
in the BERT or BERT architecture based models.
|
||||
Original paper: https://arxiv.org/abs/1606.08415
|
||||
Input shape:
|
||||
Arbitrary. Use the keyword argument `input_shape`
|
||||
(tuple of integers, does not include the samples axis)
|
||||
when using this layer as the first layer in a model.
|
||||
Output shape:
|
||||
Same shape as the input.
|
||||
|
||||
def __init__(self, approximate=True, **kwargs):
|
||||
super(GeLU, self).__init__(**kwargs)
|
||||
self.approximate = approximate
|
||||
self.supports_masking = True
|
||||
|
||||
def call(self, inputs):
|
||||
cdf = 0.5 * (1.0 + K.tanh((np.sqrt(2 / np.pi) * (inputs + 0.044715 * K.pow(inputs, 3)))))
|
||||
return inputs * cdf
|
||||
|
||||
def get_config(self):
|
||||
config = {'approximate': self.approximate}
|
||||
base_config = super(GeLU, self).get_config()
|
||||
return dict(list(base_config.items()) + list(config.items()))
|
||||
|
||||
def compute_output_shape(self, input_shape):
|
||||
return input_shape
|
||||
nn.GeLU = GeLU
|
||||
"""
|
||||
|
|
@ -33,4 +33,5 @@ def get_screen_size():
|
|||
elif 'linux' in sys.platform:
|
||||
pass
|
||||
|
||||
return (1366, 768)
|
||||
return (1366, 768)
|
||||
|
||||
Loading…
Add table
Add a link
Reference in a new issue