Repository: coreylowman/dfdx (language: rust)

Issue:
Upscale2D or Conv2D produces NaN values
My application using dfdx: https://github.com/nkoppel/generals-bot is currently failing to train a unet to play generals.io because the [policy network](https://github.com/nkoppel/generals-bot/blob/072bb025c4ce67b5e5756fa0878bcb4f8677c8e5/src/bots/ppgnet.rs#L80) produces NaN values during training. This is probably a bug in Conv2D or Upscale2D, as the NaN values only occupy a small area of the output, and this program has successfully trained networks without Upscale2D. The NaN values occur about every 100-500 batches, and appear only in the last image of a batch, in the following shape:

```
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 NaN
0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 NaN
NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN
``` 
(along the bottom and the bottom 3 of the right edge, for all 4 output channels)

Due to its relative rarity and the consistency of the shape of the NaN values, I suspect that this is a bug due to Conv2D or Upscale2D reading uninitialized memory or failing to write to uninitialized memory, but I have been unable to locate the bug thus far. It is probably more likely to be a bug in Upscale2D, because I have tried using the new cudnn implementation of Conv2D, and experienced the exact same behavior.

This error is reproduced here: https://gist.github.com/nkoppel/db07f684222a7169d2c39b205542d4c5. I have no idea why, but using Sgd or inspecting the value of `out` both prevent the bug from occurring.


Discussion / hints:
Is this with `cudnn` or `cuda`?

Edit: nvm i see it happens with both
Do you happen to have the image sizes that upscale is scaling to? I'm wondering if 
```c++
float h_scale = ((float)op.h_in-1)/(op.h_out-1);
float w_scale = ((float)op.w_in-1)/(op.w_out-1);
```
is producing nans, but that would require h_out/w_out to be `1`

Edit: but I guess that would result in ALL the values in the image being Nan, not just the edge
I suspect that this bug is memory-related, and it being difficult to reproduce is due to the pseudo-randomness in cuda's malloc algorithm. 
Also, this happens in the forward call right? Not backward?
Yes, at least within my application, it happens in the forward call. If it had happened in the backward call, parameters would have been set to NaN and caused the entire output to be NaN.
The weird thing is since upscale2d calls `let mut out = inp.device.try_zeros_like(&(chan, out_height, out_width))?;`, this will memset all the data in the output to 0. Meaning the kernel is definitely setting the nan values.
Yes, and it seems impossible for upscale2d to read out of bounds as long as the data in `op` and `inp_strides` is correct.

Fix patch:
diff --git a/src/tensor_ops/upscale2d/cuda_kernel.rs b/src/tensor_ops/upscale2d/cuda_kernel.rs
index ea5ab70e2..2d2013617 100644
--- a/src/tensor_ops/upscale2d/cuda_kernel.rs
+++ b/src/tensor_ops/upscale2d/cuda_kernel.rs
@@ -1,11 +1,11 @@
 use crate::{
     shapes::*,
-    tensor::{Cuda, Tensor},
+    tensor::{launch_cfg, Cuda, Tensor},
 };
 
 use std::sync::Arc;
 
-use cudarc::driver::{DeviceRepr, LaunchAsync, LaunchConfig};
+use cudarc::driver::{DeviceRepr, LaunchAsync};
 
 use super::{Bilinear, NearestNeighbor, UpscaleMethod};
 
@@ -13,9 +13,9 @@ const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/upscale2d.ptx"));
 
 unsafe impl DeviceRepr for super::Upscale2DOp {}
 
-fn make_4d<S: Shape>(strides: S::Concrete, pad: usize) -> [usize; 4] {
+fn make_4d<S: Shape>(strides: S::Concrete) -> [usize; 4] {
     match S::NUM_DIMS {
-        3 => [pad, strides[0], strides[1], strides[2]],
+        3 => [0, strides[0], strides[1], strides[2]],
         4 => [strides[0], strides[1], strides[2], strides[3]],
         _ => panic!("Only implemented for 3d & 4d arrays"),
     }
@@ -56,16 +56,14 @@ where
                 .load_ptx(PTX_SRC.into(), Self::FWD, &[Self::FWD, Self::BWD])?;
         }
 
-        let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides, 0).into())?;
-        let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides, 0).into())?;
+        let strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
         let fwd_fn = self.dev.get_func(Self::FWD, Self::FWD).unwrap();
-        let cfg = LaunchConfig::for_num_elems(out.shape().num_elements() as u32);
+        let cfg = launch_cfg(out.shape().num_elements() as u32);
         let params = (
-            op,                           // const Pool2dOp op,
-            &inp_strides,                 // const size_t *inp_strides,
-            &out_strides,                 // const size_t *out_strides,
-            inp.data.as_ref(),            // const float *inp,
-            Arc::make_mut(&mut out.data), // float *out
+            op,
+            &strides,
+            inp.data.as_ref(),
+            Arc::make_mut(&mut out.data),
         );
         unsafe { fwd_fn.launch(cfg, params) }?;
         Ok(())
@@ -78,17 +76,10 @@ where
         out: &Tensor<O, E, Self>,
         grad_out: &Self::Vec<E>,
     ) -> Result<(), Self::Err> {
-        let inp_strides = self.dev.htod_copy(make_4d::<I>(inp.strides, 0).into())?;
-        let out_strides = self.dev.htod_copy(make_4d::<O>(out.strides, 0).into())?;
+        let strides = self.dev.htod_copy(make_4d::<I>(inp.strides).into())?;
         let bwd_fn = self.dev.get_func(Self::FWD, Self::BWD).unwrap();
-        let cfg = LaunchConfig::for_num_elems(out.shape().num_elements() as u32);
-        let params = (
-            op,           // const Pool2dOp op,
-            &inp_strides, // const size_t *inp_strides,
-            &out_strides, // const size_t *out_strides,
-            grad_inp,     // float *grad_inp,
-            grad_out,     // const float *grad_out
-        );
+        let cfg = launch_cfg(out.shape().num_elements() as u32);
+        let params = (op, &strides, grad_inp, grad_out);
         unsafe { bwd_fn.launch(cfg, params) }?;
         Ok(())
     }
diff --git a/src/tensor_ops/upscale2d/upscale2d.cu b/src/tensor_ops/upscale2d/upscale2d.cu
index 39786f4e5..49ad1b91a 100644
--- a/src/tensor_ops/upscale2d/upscale2d.cu
+++ b/src/tensor_ops/upscale2d/upscale2d.cu
@@ -13,7 +13,6 @@ template<typename T>
 __device__ void nearest_upscale2d_fwd(
     const Upscale2dOp op,
     const size_t *inp_strides,
-    const size_t *out_strides,
     const T *inp, // 4d (Batch, Channels, Height, Width)
     T *out // 4d (Batch, Channels, HeightOut, WidthOut)
 ) {
@@ -34,8 +33,8 @@ __device__ void nearest_upscale2d_fwd(
     idx /= op.chan;
     const size_t b = idx % op.batch;
 
-    size_t ih = min(static_cast<size_t>(h_scale * oh), op.h_out - 1);
-    size_t iw = min(static_cast<size_t>(w_scale * ow), op.w_out - 1);
+    size_t ih = min(static_cast<size_t>(h_scale * oh), op.h_in - 1);
+    size_t iw = min(static_cast<size_t>(w_scale * ow), op.w_in - 1);
 
     size_t inp_i = b * inp_strides[0] + c * inp_strides[1] + ih * inp_strides[2] + iw * inp_strides[3];
     
@@ -46,7 +45,6 @@ template<typename T>
 __device__ void nearest_upscale2d_bwd(
     const Upscale2dOp op,
     const size_t *inp_strides,
-    const size_t *out_strides,
     T *grad_inp,
     const T *grad_out // 4d (Batch, Channels, HeightOut, WidthOut)
 ) {
@@ -67,8 +65,8 @@ __device__ void nearest_upscale2d_bwd(
     idx /= op.chan;
     const size_t b = idx % op.batch;
 
-    size_t ih = min(static_cast<size_t>(h_scale * oh), op.h_out - 1);
-    size_t iw = min(static_cast<size_t>(w_scale * ow), op.w_out - 1);
+    size_t ih = min(static_cast<size_t>(h_scale * oh), op.h_in - 1);
+    size_t iw = min(static_cast<size_t>(w_scale * ow), op.w_in - 1);
 
     size_t inp_i = b * inp_strides[0] + c * inp_strides[1] + ih * inp_strides[2] + iw * inp_strides[3];
     atomicAdd(grad_inp + inp_i, grad_out[i]);
@@ -78,7 +76,6 @@ template<typename T>
 __device__ void bilinear_upscale2d_fwd(
     const Upscale2dOp op,
     const size_t *inp_strides,
-    const size_t *out_strides,
     const T *inp, // 4d (Batch, Channels, Height, Width)
     T *out // 4d (Batch, Channels, HeightOut, WidthOut)
 ) {
@@ -98,12 +95,11 @@ __device__ void bilinear_upscale2d_fwd(
     const size_t c = idx % op.chan;
     idx /= op.chan;
     const size_t b = idx % op.batch;
-    idx /= op.batch;
 
-    size_t y0 = min(static_cast<size_t>(h_scale * oh), op.h_out - 1);
-    size_t y1 = min(y0 + 1, op.h_out - 1);
-    size_t x0 = min(static_cast<size_t>(w_scale * ow), op.w_out - 1);
-    size_t x1 = min(x0 + 1, op.w_out - 1);
+    size_t y0 = min(static_cast<size_t>(h_scale * oh), op.h_in - 1);
+    size_t y1 = min(y0 + 1, op.h_in - 1);
+    size_t x0 = min(static_cast<size_t>(w_scale * ow), op.w_in - 1);
+    size_t x1 = min(x0 + 1, op.w_in - 1);
 
     T hs = h_scale * oh - y0;
     T ws = w_scale * ow - x0;
@@ -122,7 +118,6 @@ template<typename T>
 __device__ void bilinear_upscale2d_bwd(
     const Upscale2dOp op,
     const size_t *inp_strides,
-    const size_t *out_strides,
     T *grad_inp, // 4d (Batch, Channels, Height, Width)
     const T *grad_out // 4d (Batch, Channels, HeightOut, WidthOut)
 ) {
@@ -142,12 +137,11 @@ __device__ void bilinear_upscale2d_bwd(
     const size_t c = idx % op.chan;
     idx /= op.chan;
     const size_t b = idx % op.batch;
-    idx /= op.batch;
 
-    size_t y0 = min(static_cast<size_t>(h_scale * oh), op.h_out - 1);
-    size_t y1 = min(y0 + 1, op.h_out - 1);
-    size_t x0 = min(static_cast<size_t>(w_scale * ow), op.w_out - 1);
-    size_t x1 = min(x0 + 1, op.w_out - 1);
+    size_t y0 = min(static_cast<size_t>(h_scale * oh), op.h_in - 1);
+    size_t y1 = min(y0 + 1, op.h_in - 1);
+    size_t x0 = min(static_cast<size_t>(w_scale * ow), op.w_in - 1);
+    size_t x1 = min(x0 + 1, op.w_in - 1);
 
     T hs = h_scale * oh - y0;
     T ws = w_scale * ow - x0;
@@ -166,20 +160,18 @@ __device__ void bilinear_upscale2d_bwd(
 extern "C" __global__ void fwd( \
     const Upscale2dOp op, \
     const size_t *inp_strides, \
-    const size_t *out_strides, \
     const TYPENAME *inp, \
     TYPENAME *out \
 ) { \
-    fwd_FN(op, inp_strides, out_strides, inp, out); \
+    fwd_FN(op, inp_strides, inp, out); \
 } \
 extern "C" __global__ void bwd( \
     const Upscale2dOp op, \
     const size_t *inp_strides, \
-    const size_t *out_strides, \
     TYPENAME *grad_inp, \
     const TYPENAME *grad_out \
 ) { \
-    bwd_FN(op, inp_strides, out_strides, grad_inp, grad_out); \
+    bwd_FN(op, inp_strides, grad_inp, grad_out); \
 }
 
 UPSCALE_OP(


Test patch: