solved BatchNorm1d problem

Author: Zhangyanbo

src/INN/INN.py

@@ -171,7 +171,7 @@ def forward(self, x, log_p=0, log_det_J=0):
                 var = self.running_var # [dim]
             else:
                 # if in training
-                var = torch.var(x, dim=0, unbiased=False) # [dim]
+                var = torch.var(x, dim=0, unbiased=False).detach() # [dim]
 
             x = super(BatchNorm1d, self).forward(x)
 

tests/quick_tests.ipynb

@@ -209,11 +209,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 14,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2021-04-25T20:16:59.907029Z",
-     "start_time": "2021-04-25T20:16:59.900497Z"
+     "end_time": "2021-04-25T20:27:57.577833Z",
+     "start_time": "2021-04-25T20:27:57.573240Z"
     }
    },
    "outputs": [],
@@ -223,21 +223,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 15,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2021-04-25T20:17:21.038359Z",
-     "start_time": "2021-04-25T20:17:21.015209Z"
+     "end_time": "2021-04-25T20:27:57.883569Z",
+     "start_time": "2021-04-25T20:27:57.873327Z"
     }
    },
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "J_g=tensor([-1.5020e-05, -1.5020e-05, -1.5020e-05, -1.5020e-05, -1.5020e-05,\n",
-      "        -1.5020e-05]),\n",
-      "J_c=-1.5020295904832892e-05\n"
+      "J_g=tensor([ 1.9498, -0.0728,  1.3574,  0.7588,  1.3956,  0.7595]),\n",
+      "J_c=tensor([ 1.9498, -0.0728,  1.3574,  0.7588,  1.3956,  0.7595])\n"
      ]
     }
    ],
@@ -251,23 +250,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 16,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2021-04-25T20:17:23.285413Z",
-     "start_time": "2021-04-25T20:17:23.279627Z"
+     "end_time": "2021-04-25T20:28:05.686690Z",
+     "start_time": "2021-04-25T20:28:05.681378Z"
     }
    },
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "tensor([[1.0000, 0.0000, 0.0000],\n",
-       "        [0.0000, 1.0000, 0.0000],\n",
-       "        [0.0000, 0.0000, 1.0000]])"
+       "tensor([[ 1.9760, -1.3065,  1.2205],\n",
+       "        [ 0.5567, -0.8171, -1.6935],\n",
+       "        [ 0.0597,  2.1638,  2.1138]])"
       ]
      },
-     "execution_count": 12,
+     "execution_count": 16,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -420,40 +419,40 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 17,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2021-04-25T20:25:00.696600Z",
-     "start_time": "2021-04-25T20:25:00.688777Z"
+     "end_time": "2021-04-25T20:51:57.518139Z",
+     "start_time": "2021-04-25T20:51:57.513012Z"
     }
    },
    "outputs": [],
    "source": [
     "x = torch.randn((5, 3))\n",
-    "bn = nn.BatchNorm1d(3)"
+    "bn = nn.BatchNorm1d(3, affine=False)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 18,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2021-04-25T20:25:01.168809Z",
-     "start_time": "2021-04-25T20:25:01.158170Z"
+     "end_time": "2021-04-25T20:51:58.028429Z",
+     "start_time": "2021-04-25T20:51:58.009478Z"
     }
    },
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "tensor([[-1.1245,  1.3747, -0.5232],\n",
-       "        [ 1.1031,  0.5880,  0.0641],\n",
-       "        [-0.8741,  0.0202, -1.3274],\n",
-       "        [-0.3750, -0.3597,  0.0675],\n",
-       "        [ 1.2705, -1.6232,  1.7191]], grad_fn=<NativeBatchNormBackward>)"
+       "tensor([[-1.6941,  0.2933, -0.2451],\n",
+       "        [-0.1313, -0.2711,  1.4740],\n",
+       "        [ 0.2754, -0.2282,  0.4445],\n",
+       "        [ 0.1287, -1.4409, -0.0721],\n",
+       "        [ 1.4213,  1.6469, -1.6014]])"
       ]
      },
-     "execution_count": 20,
+     "execution_count": 18,
      "metadata": {},
      "output_type": "execute_result"
     }