init

2022-05-25 19:54:48 +02:00
parent b8b42736f3
commit b098ecee9c
67 changed files with 39844 additions and 1 deletions
@@ -0,0 +1,146 @@
+import os
+import numpy as np
+import tensorflow as tf
+from stylegan2.layers.cuda import custom_ops
+
+
+def _get_plugin():
+    loc = os.path.dirname(os.path.abspath(__file__))
+    cu_fn = 'upfirdn_2d.cu'
+    return custom_ops.get_plugin(os.path.join(loc, cu_fn))
+
+
+def _setup_kernel(k):
+    k = np.asarray(k, dtype=np.float32)
+    if k.ndim == 1:
+        k = np.outer(k, k)
+    k /= np.sum(k)
+    assert k.ndim == 2
+    assert k.shape[0] == k.shape[1]
+    return k
+
+
+def compute_paddings(resample_kernel, convW, up, down, is_conv, factor=2, gain=1):
+    assert not (up and down)
+
+    k = [1] * factor if resample_kernel is None else resample_kernel
+    if up:
+        k = _setup_kernel(k) * (gain * (factor ** 2))
+        if is_conv:
+            p = (k.shape[0] - factor) - (convW - 1)
+            pad0 = (p + 1) // 2 + factor - 1
+            pad1 = p // 2 + 1
+        else:
+            p = k.shape[0] - factor
+            pad0 = (p + 1) // 2 + factor - 1
+            pad1 = p // 2
+    elif down:
+        k = _setup_kernel(k) * gain
+        if is_conv:
+            p = (k.shape[0] - factor) + (convW - 1)
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+        else:
+            p = k.shape[0] - factor
+            pad0 = (p + 1) // 2
+            pad1 = p // 2
+    else:
+        k = resample_kernel
+        pad0, pad1 = 0, 0
+    return k, pad0, pad1
+
+
+def upsample_2d(x, x_res, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+    return _simple_upfirdn_2d(x, x_res, k, up=factor, pad0=pad0, pad1=pad1)
+
+
+def downsample_2d(x, x_res, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+    return _simple_upfirdn_2d(x, x_res, k, down=factor, pad0=pad0, pad1=pad1)
+
+
+def upsample_conv_2d(x, x_res, w, convH, convW, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+
+    # Check weight shape.
+    w = tf.convert_to_tensor(w)
+    assert w.shape.rank == 4
+    # convH = w.shape[0]
+    # convW = w.shape[1]
+    inC = tf.shape(w)[2]
+    outC = tf.shape(w)[3]
+    assert convW == convH
+
+    # Determine data dimensions.
+    stride = [1, 1, factor, factor]
+    output_shape = [tf.shape(x)[0], outC, (x_res - 1) * factor + convH, (int(x_res/2) - 1) * factor + convW]
+    num_groups = tf.shape(x)[1] // inC
+
+    # Transpose weights.
+    w = tf.reshape(w, [convH, convW, inC, num_groups, -1])
+    w = tf.transpose(w[::-1, ::-1], [0, 1, 4, 3, 2])
+    w = tf.reshape(w, [convH, convW, -1, num_groups * inC])
+
+    # Execute.
+    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format='NCHW')
+    new_x_res = output_shape[2]
+    return _simple_upfirdn_2d(x, new_x_res, k, pad0=pad0, pad1=pad1)
+
+
+def conv_downsample_2d(x, x_res, w, convH, convW, pad0, pad1, k, factor=2):
+    assert isinstance(factor, int) and factor >= 1
+    w = tf.convert_to_tensor(w)
+    # convH, convW, _inC, _outC = w.shape.as_list()
+    assert convW == convH
+
+    s = [1, 1, factor, factor]
+    x = _simple_upfirdn_2d(x, x_res, k, pad0=pad0, pad1=pad1)
+    return tf.nn.conv2d(x, w, strides=s, padding='VALID', data_format='NCHW')
+
+
+def _simple_upfirdn_2d(x, x_res, k, up=1, down=1, pad0=0, pad1=0):
+    assert x.shape.rank == 4
+    y = x
+    y = tf.reshape(y, [-1, tf.shape(x)[2], tf.shape(x)[3], 1])
+    y = upfirdn_2d_cuda(y, k, upx=up, upy=up, downx=down, downy=down, padx0=pad0, padx1=pad1, pady0=pad0, pady1=pad1)
+    y = tf.reshape(y, [-1, tf.shape(x)[1], tf.shape(y)[1], tf.shape(y)[2]])
+    return y
+
+
+def upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
+    """Fast CUDA implementation of `upfirdn_2d()` using custom ops."""
+
+    x = tf.convert_to_tensor(x)
+    k = np.asarray(k, dtype=np.float32)
+    majorDim, inH, inW, minorDim = x.shape.as_list()
+    kernelH, kernelW = k.shape
+    assert inW >= 1 and inH >= 1
+    assert kernelW >= 1 and kernelH >= 1
+    assert isinstance(upx, int) and isinstance(upy, int)
+    assert isinstance(downx, int) and isinstance(downy, int)
+    assert isinstance(padx0, int) and isinstance(padx1, int)
+    assert isinstance(pady0, int) and isinstance(pady1, int)
+
+    outW = (inW * upx + padx0 + padx1 - kernelW) // downx + 1
+    outH = (inH * upy + pady0 + pady1 - kernelH) // downy + 1
+    assert outW >= 1 and outH >= 1
+
+    kc = tf.constant(k, dtype=x.dtype)
+    gkc = tf.constant(k[::-1, ::-1], dtype=x.dtype)
+    gpadx0 = kernelW - padx0 - 1
+    gpady0 = kernelH - pady0 - 1
+    gpadx1 = inW * upx - outW * downx + padx0 - upx + 1
+    gpady1 = inH * upy - outH * downy + pady0 - upy + 1
+
+    @tf.custom_gradient
+    def func(x):
+        y = _get_plugin().up_fir_dn2d(x=x, k=kc, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)
+        y.set_shape([majorDim, outH, outW, minorDim])
+        @tf.custom_gradient
+        def grad(dy):
+            dx = _get_plugin().up_fir_dn2d(x=dy, k=gkc, upx=downx, upy=downy, downx=upx, downy=upy, padx0=gpadx0, padx1=gpadx1, pady0=gpady0, pady1=gpady1)
+            dx.set_shape([majorDim, inH, inW, minorDim])
+            return dx, func
+        return y, grad
+    return func(x)