impl & test for all case by adding new initer module

Author: HugoPhi

cnn_cifar10.ipynb

@@ -129,38 +129,24 @@
     "from plugins.minitorch.optimizer import Adam\n",
     "from plugins.minitorch.utils import softmax, cross_entropy_loss\n",
     "from plugins.minitorch.initer import Initer\n",
-    "from plugins.minitorch.nn import dropout\n",
+    "from plugins.minitorch.nn import Dense\n",
+    "from plugins.minitorch.loss import CrossEntropyLoss\n",
     "\n",
     "class mlp_clf:\n",
     "    def __init__(self, lr=0.01):\n",
     "        self.config = {\n",
-    "            'fc4relu:0' : {\n",
-    "                'input_dim': 32 * 32 * 3,\n",
-    "                'output_dim': 128 * 3,\n",
-    "            },\n",
-    "            'relu': {},\n",
-    "            'fc4relu:1' : {\n",
-    "                'input_dim': 128 * 3,\n",
-    "                'output_dim': 64 * 3,\n",
-    "            },\n",
-    "            'fc4relu:2' : {\n",
-    "                'input_dim': 64 * 3,\n",
-    "                'output_dim': 32 * 3,\n",
-    "            },\n",
-    "            'fc4relu:3' : {\n",
-    "                'input_dim': 32 * 3,\n",
-    "                'output_dim': 16 * 3,\n",
-    "            },\n",
-    "            'fc4relu:4' : {\n",
-    "                'input_dim': 16 * 3,\n",
-    "                'output_dim': 10,\n",
-    "            }\n",
+    "            'fc:0': Dense.get_linear(32 * 32 * 3, 128 * 3),\n",
+    "            'fc:1': Dense.get_linear(128 * 3, 64 * 3),\n",
+    "            'fc:2': Dense.get_linear(64 * 3, 32 * 3),\n",
+    "            'fc:3': Dense.get_linear(32 * 3, 16 * 3),\n",
+    "            'fc:4': Dense.get_linear(16 * 3, 10),\n",
     "        }\n",
     "\n",
     "        initer = Initer(self.config, key)\n",
     "        self.optr = Adam(initer(), lr=lr, batch_size=512)\n",
+    "        self.losser = CrossEntropyLoss(self.forward)\n",
     "\n",
-    "    def predict_proba(self, x: jnp.ndarray, params, train=False):\n",
+    "    def forward(self, x: jnp.ndarray, params, train=False):\n",
     "        res = x\n",
     "        key = random.PRNGKey(42)\n",
     "        for p in params.values():\n",
@@ -172,42 +158,37 @@
     "            # 后面发现即使不使用JIT也会出现相同的问题，因为没有JIT也会有其它多线程的优化，导致这里发生问题。\n",
     "            res = res @ p['w'] + p['b']  \n",
     "            res = jnp.maximum(0, res)  # use relu activation function\n",
-    "            res, key = dropout(res, key, p=0.1, train=train)\n",
+    "            res, key = Dense.dropout(res, key, p=0.1, train=train)\n",
     "\n",
     "        return softmax(res)\n",
     "\n",
     "    def fit(self, x_train, y_train_proba, x_test, y_test_proba, epoches=100): \n",
-    "        cnt = 0\n",
-    "\n",
     "        @jit\n",
     "        def _acc(y_true_proba, y_pred_proba):\n",
     "            y_true = jnp.argmax(y_true_proba, axis=1)\n",
     "            y_pred = jnp.argmax(y_pred_proba, axis=1)\n",
     "            return jnp.mean(y_true == y_pred)\n",
     "\n",
-    "        _loss = lambda params, x, y_true: cross_entropy_loss(y_true, self.predict_proba(x, params, True)) \n",
-    "        _loss = jit(_loss)  # accelerate loss function by JIT\n",
+    "        _loss = self.losser.get_loss(train=True)\n",
+    "        _loss = jit(_loss)\n",
     "        self.optr.open(_loss, x_train, y_train_proba)\n",
     "        \n",
-    "        _tloss = lambda params: cross_entropy_loss(y_test_proba, self.predict_proba(x_test, params, False)) \n",
-    "        _tloss = jit(_tloss)  # accelerate loss function by JIT\n",
+    "        _tloss = self.losser.get_embed_loss(x_test, y_test_proba, train=False)\n",
+    "        _tloss = jit(_tloss)\n",
     "        \n",
     "\n",
     "        acc, loss, tacc, tloss = [], [], [], []  # train acc, train loss, test acc, test loss\n",
     "\n",
-    "        for _ in range(epoches):\n",
+    "        for cnt in range(epoches):\n",
     "            loss.append(_loss(self.optr.get_params(), x_train, y_train_proba))\n",
     "            tloss.append(_tloss(self.optr.get_params()))\n",
     "\n",
-    "            self.train = True  # use dropout only while updating grads\n",
     "            self.optr.update()\n",
-    "            self.train = False\n",
     "            \n",
-    "            acc.append(_acc(y_train_proba, self.predict_proba(x_train, self.optr.get_params())))\n",
-    "            tacc.append(_acc(y_test_proba, self.predict_proba(x_test, self.optr.get_params())))\n",
-    "            cnt += 1\n",
-    "            if cnt % 10 == 0:\n",
-    "                print(f'>> epoch: {cnt}, train acc: {acc[-1]}, test acc: {tacc[-1]}')\n",
+    "            acc.append(_acc(y_train_proba, self.forward(x_train, self.optr.get_params())))\n",
+    "            tacc.append(_acc(y_test_proba, self.forward(x_test, self.optr.get_params())))\n",
+    "            if (cnt + 1) % 10 == 0:\n",
+    "                print(f'>> epoch: {cnt + 1}, train acc: {acc[-1]}, test acc: {tacc[-1]}')\n",
     "\n",
     "        return acc, loss, tacc, tloss"
    ]