File size: 7,355 Bytes
2fe3da0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
/*
* Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
*
* NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
* property and proprietary rights in and to this material, related
* documentation and any modifications thereto. Any use, reproduction,
* disclosure or distribution of this material and related documentation
* without an express license agreement from NVIDIA CORPORATION or
* its affiliates is strictly prohibited.
*/
#include <cuda.h>
#include "common.h"
#include "loss.h"
//------------------------------------------------------------------------
// Utils
__device__ inline float bwdAbs(float x) { return x == 0.0f ? 0.0f : x < 0.0f ? -1.0f : 1.0f; }
__device__ float warpSum(float val) {
for (int i = 1; i < 32; i *= 2)
val += __shfl_xor_sync(0xFFFFFFFF, val, i);
return val;
}
//------------------------------------------------------------------------
// Tonemapping
__device__ inline float fwdSRGB(float x)
{
return x > 0.0031308f ? powf(max(x, 0.0031308f), 1.0f / 2.4f) * 1.055f - 0.055f : 12.92f * max(x, 0.0f);
}
__device__ inline void bwdSRGB(float x, float &d_x, float d_out)
{
if (x > 0.0031308f)
d_x += d_out * 0.439583f / powf(x, 0.583333f);
else if (x > 0.0f)
d_x += d_out * 12.92f;
}
__device__ inline vec3f fwdTonemapLogSRGB(vec3f x)
{
return vec3f(fwdSRGB(logf(x.x + 1.0f)), fwdSRGB(logf(x.y + 1.0f)), fwdSRGB(logf(x.z + 1.0f)));
}
__device__ inline void bwdTonemapLogSRGB(vec3f x, vec3f& d_x, vec3f d_out)
{
if (x.x > 0.0f && x.x < 65535.0f)
{
bwdSRGB(logf(x.x + 1.0f), d_x.x, d_out.x);
d_x.x *= 1 / (x.x + 1.0f);
}
if (x.y > 0.0f && x.y < 65535.0f)
{
bwdSRGB(logf(x.y + 1.0f), d_x.y, d_out.y);
d_x.y *= 1 / (x.y + 1.0f);
}
if (x.z > 0.0f && x.z < 65535.0f)
{
bwdSRGB(logf(x.z + 1.0f), d_x.z, d_out.z);
d_x.z *= 1 / (x.z + 1.0f);
}
}
__device__ inline float fwdRELMSE(float img, float target, float eps = 0.1f)
{
return (img - target) * (img - target) / (img * img + target * target + eps);
}
__device__ inline void bwdRELMSE(float img, float target, float &d_img, float &d_target, float d_out, float eps = 0.1f)
{
float denom = (target * target + img * img + eps);
d_img += d_out * 2 * (img - target) * (target * (target + img) + eps) / (denom * denom);
d_target -= d_out * 2 * (img - target) * (img * (target + img) + eps) / (denom * denom);
}
__device__ inline float fwdSMAPE(float img, float target, float eps=0.01f)
{
return abs(img - target) / (img + target + eps);
}
__device__ inline void bwdSMAPE(float img, float target, float& d_img, float& d_target, float d_out, float eps = 0.01f)
{
float denom = (target + img + eps);
d_img += d_out * bwdAbs(img - target) * (2 * target + eps) / (denom * denom);
d_target -= d_out * bwdAbs(img - target) * (2 * img + eps) / (denom * denom);
}
//------------------------------------------------------------------------
// Kernels
__global__ void imgLossFwdKernel(LossKernelParams p)
{
// Calculate pixel position.
unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
unsigned int pz = blockIdx.z;
float floss = 0.0f;
if (px < p.gridSize.x && py < p.gridSize.y && pz < p.gridSize.z)
{
vec3f img = p.img.fetch3(px, py, pz);
vec3f target = p.target.fetch3(px, py, pz);
img = vec3f(clamp(img.x, 0.0f, 65535.0f), clamp(img.y, 0.0f, 65535.0f), clamp(img.z, 0.0f, 65535.0f));
target = vec3f(clamp(target.x, 0.0f, 65535.0f), clamp(target.y, 0.0f, 65535.0f), clamp(target.z, 0.0f, 65535.0f));
if (p.tonemapper == TONEMAPPER_LOG_SRGB)
{
img = fwdTonemapLogSRGB(img);
target = fwdTonemapLogSRGB(target);
}
vec3f vloss(0);
if (p.loss == LOSS_MSE)
vloss = (img - target) * (img - target);
else if (p.loss == LOSS_RELMSE)
vloss = vec3f(fwdRELMSE(img.x, target.x), fwdRELMSE(img.y, target.y), fwdRELMSE(img.z, target.z));
else if (p.loss == LOSS_SMAPE)
vloss = vec3f(fwdSMAPE(img.x, target.x), fwdSMAPE(img.y, target.y), fwdSMAPE(img.z, target.z));
else
vloss = vec3f(abs(img.x - target.x), abs(img.y - target.y), abs(img.z - target.z));
floss = sum(vloss) / 3.0f;
}
floss = warpSum(floss);
dim3 warpSize = getWarpSize(blockDim);
if (px < p.gridSize.x && py < p.gridSize.y && pz < p.gridSize.z && threadIdx.x % warpSize.x == 0 && threadIdx.y % warpSize.y == 0 && threadIdx.z % warpSize.z == 0)
p.out.store(px / warpSize.x, py / warpSize.y, pz / warpSize.z, floss);
}
__global__ void imgLossBwdKernel(LossKernelParams p)
{
// Calculate pixel position.
unsigned int px = blockIdx.x * blockDim.x + threadIdx.x;
unsigned int py = blockIdx.y * blockDim.y + threadIdx.y;
unsigned int pz = blockIdx.z;
if (px >= p.gridSize.x || py >= p.gridSize.y || pz >= p.gridSize.z)
return;
dim3 warpSize = getWarpSize(blockDim);
vec3f _img = p.img.fetch3(px, py, pz);
vec3f _target = p.target.fetch3(px, py, pz);
float d_out = p.out.fetch1(px / warpSize.x, py / warpSize.y, pz / warpSize.z);
/////////////////////////////////////////////////////////////////////
// FWD
vec3f img = _img, target = _target;
if (p.tonemapper == TONEMAPPER_LOG_SRGB)
{
img = fwdTonemapLogSRGB(img);
target = fwdTonemapLogSRGB(target);
}
/////////////////////////////////////////////////////////////////////
// BWD
vec3f d_vloss = vec3f(d_out, d_out, d_out) / 3.0f;
vec3f d_img(0), d_target(0);
if (p.loss == LOSS_MSE)
{
d_img = vec3f(d_vloss.x * 2 * (img.x - target.x), d_vloss.y * 2 * (img.y - target.y), d_vloss.x * 2 * (img.z - target.z));
d_target = -d_img;
}
else if (p.loss == LOSS_RELMSE)
{
bwdRELMSE(img.x, target.x, d_img.x, d_target.x, d_vloss.x);
bwdRELMSE(img.y, target.y, d_img.y, d_target.y, d_vloss.y);
bwdRELMSE(img.z, target.z, d_img.z, d_target.z, d_vloss.z);
}
else if (p.loss == LOSS_SMAPE)
{
bwdSMAPE(img.x, target.x, d_img.x, d_target.x, d_vloss.x);
bwdSMAPE(img.y, target.y, d_img.y, d_target.y, d_vloss.y);
bwdSMAPE(img.z, target.z, d_img.z, d_target.z, d_vloss.z);
}
else
{
d_img = d_vloss * vec3f(bwdAbs(img.x - target.x), bwdAbs(img.y - target.y), bwdAbs(img.z - target.z));
d_target = -d_img;
}
if (p.tonemapper == TONEMAPPER_LOG_SRGB)
{
vec3f d__img(0), d__target(0);
bwdTonemapLogSRGB(_img, d__img, d_img);
bwdTonemapLogSRGB(_target, d__target, d_target);
d_img = d__img; d_target = d__target;
}
if (_img.x <= 0.0f || _img.x >= 65535.0f) d_img.x = 0;
if (_img.y <= 0.0f || _img.y >= 65535.0f) d_img.y = 0;
if (_img.z <= 0.0f || _img.z >= 65535.0f) d_img.z = 0;
if (_target.x <= 0.0f || _target.x >= 65535.0f) d_target.x = 0;
if (_target.y <= 0.0f || _target.y >= 65535.0f) d_target.y = 0;
if (_target.z <= 0.0f || _target.z >= 65535.0f) d_target.z = 0;
p.img.store_grad(px, py, pz, d_img);
p.target.store_grad(px, py, pz, d_target);
} |