Repository: IBM/sail (language: python)

Issue:
Omni Scale 1D CNN implementation
To convert OS 1D CNN into a scikit-learn compatible neural network library that's wraped with PyTorch.


Discussion / hints:

Fix patch:
diff --git a/notebooks/os_cnn_classification.ipynb b/notebooks/os_cnn_classification.ipynb
new file mode 100644
index 0000000..61bce0b
--- /dev/null
+++ b/notebooks/os_cnn_classification.ipynb
@@ -0,0 +1,291 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "f371e50e-fb5c-4f5f-a39a-4258b7deb6e5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/abduvoris/anaconda3/envs/imla_assignment/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import sys\n",
+    "import os\n",
+    "os.chdir(\"..\")\n",
+    "\n",
+    "from os.path import dirname\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader, TensorDataset\n",
+    "from sklearn import preprocessing\n",
+    "\n",
+    "from sail.models.torch.os_cnn import OS_CNN_CLassifier"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "17fb7c93-7d34-4587-b7e2-89ad6f20cb5f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# methods for preprocessing data \n",
+    "def set_nan_to_zero(a):\n",
+    "    where_are_NaNs = np.isnan(a)\n",
+    "    a[where_are_NaNs] = 0\n",
+    "    return a\n",
+    "\n",
+    "def TSC_data_loader(dataset_path,dataset_name):\n",
+    "    Train_dataset = np.loadtxt(\n",
+    "        dataset_path + '/' + dataset_name + '/' + dataset_name + '_TRAIN.tsv')\n",
+    "    Test_dataset = np.loadtxt(\n",
+    "        dataset_path + '/' + dataset_name + '/' + dataset_name + '_TEST.tsv')\n",
+    "    Train_dataset = Train_dataset.astype(np.float32)\n",
+    "    Test_dataset = Test_dataset.astype(np.float32)\n",
+    "\n",
+    "    X_train = Train_dataset[:, 1:]\n",
+    "    y_train = Train_dataset[:, 0:1]\n",
+    "\n",
+    "    X_test = Test_dataset[:, 1:]\n",
+    "    y_test = Test_dataset[:, 0:1]\n",
+    "    le = preprocessing.LabelEncoder()\n",
+    "    le.fit(np.squeeze(y_train, axis=1))\n",
+    "    y_train = le.transform(np.squeeze(y_train, axis=1))\n",
+    "    y_test = le.transform(np.squeeze(y_test, axis=1))\n",
+    "    return set_nan_to_zero(X_train), y_train, set_nan_to_zero(X_test), y_test\n",
+    "\n",
+    "\n",
+    "def load_to_torch(X_train, y_train, X_test, y_test, device):\n",
+    "    X_train = torch.from_numpy(X_train)\n",
+    "    X_train.requires_grad = False\n",
+    "    X_train = X_train.to(device)\n",
+    "    y_train = torch.from_numpy(y_train).to(device)\n",
+    "\n",
+    "    X_test = torch.from_numpy(X_test)\n",
+    "    X_test.requires_grad = False\n",
+    "    X_test = X_test.to(device)\n",
+    "    y_test = torch.from_numpy(y_test).to(device)\n",
+    "\n",
+    "\n",
+    "    if len(X_train.shape) == 2:\n",
+    "        X_train = X_train.unsqueeze_(1)\n",
+    "        X_test = X_test.unsqueeze_(1)\n",
+    "    return X_train, y_train, X_test, y_test"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "e80cc71e-edc6-4df2-917f-bfa20f20d004",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "train data shape (450, 270)\n",
+      "\n",
+      "train label shape (450,)\n",
+      "test data shape (455, 270)\n",
+      "test label shape (455,)\n",
+      "unique train label [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23\n",
+      " 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47\n",
+      " 48 49]\n",
+      "unique test label [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23\n",
+      " 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47\n",
+      " 48 49]\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/abduvoris/anaconda3/envs/imla_assignment/lib/python3.9/site-packages/sklearn/model_selection/_split.py:680: UserWarning: The least populated class in y has only 1 members, which is less than n_splits=5.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "  epoch    train_loss    valid_acc    valid_loss     dur\n",
+      "-------  ------------  -----------  ------------  ------\n",
+      "      1        \u001b[36m3.7022\u001b[0m       \u001b[32m0.2222\u001b[0m        \u001b[35m3.3004\u001b[0m  1.4693\n",
+      "      2        \u001b[36m3.1146\u001b[0m       \u001b[32m0.2333\u001b[0m        \u001b[35m3.0906\u001b[0m  1.4875\n",
+      "      3        \u001b[36m2.7727\u001b[0m       \u001b[32m0.2556\u001b[0m        \u001b[35m2.6256\u001b[0m  1.4676\n",
+      "      4        \u001b[36m2.4804\u001b[0m       \u001b[32m0.3111\u001b[0m        \u001b[35m2.3502\u001b[0m  1.4532\n",
+      "      5        \u001b[36m2.2411\u001b[0m       \u001b[32m0.3889\u001b[0m        \u001b[35m2.0951\u001b[0m  1.4549\n",
+      "      6        \u001b[36m2.0516\u001b[0m       0.3778        2.1018  1.4673\n",
+      "      7        \u001b[36m1.8672\u001b[0m       \u001b[32m0.4556\u001b[0m        \u001b[35m1.9292\u001b[0m  1.4843\n",
+      "      8        \u001b[36m1.6836\u001b[0m       \u001b[32m0.4778\u001b[0m        \u001b[35m1.8325\u001b[0m  1.5251\n",
+      "      9        \u001b[36m1.5213\u001b[0m       0.4444        1.9530  1.5033\n",
+      "     10        \u001b[36m1.3636\u001b[0m       \u001b[32m0.5889\u001b[0m        \u001b[35m1.5287\u001b[0m  1.6161\n",
+      "     11        \u001b[36m1.1743\u001b[0m       0.5111        1.6169  1.5336\n",
+      "     12        \u001b[36m1.0068\u001b[0m       0.5667        1.5536  1.5055\n",
+      "     13        \u001b[36m0.9012\u001b[0m       \u001b[32m0.6333\u001b[0m        \u001b[35m1.4182\u001b[0m  1.5172\n",
+      "     14        \u001b[36m0.7808\u001b[0m       \u001b[32m0.6889\u001b[0m        \u001b[35m1.3519\u001b[0m  1.5905\n",
+      "     15        \u001b[36m0.6509\u001b[0m       0.6667        1.3796  1.5775\n",
+      "     16        \u001b[36m0.5332\u001b[0m       \u001b[32m0.7444\u001b[0m        \u001b[35m1.2529\u001b[0m  1.5590\n",
+      "     17        \u001b[36m0.4389\u001b[0m       0.7111        \u001b[35m1.2458\u001b[0m  1.5493\n",
+      "     18        \u001b[36m0.3755\u001b[0m       0.6889        1.3482  1.5447\n",
+      "     19        \u001b[36m0.3099\u001b[0m       0.6556        1.3700  2.0928\n",
+      "     20        \u001b[36m0.2855\u001b[0m       0.7111        1.3693  2.0708\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<class 'sail.models.torch.os_cnn.OS_CNN_CLassifier'>[initialized](\n",
+       "  module_=_OS_CNN(\n",
+       "    (net): Sequential(\n",
+       "      (0): build_layer_with_layer_parameter(\n",
+       "        (padding): ConstantPad1d(padding=(33, 33), value=0)\n",
+       "        (conv1d): Conv1d(1, 20, kernel_size=(67,), stride=(1,))\n",
+       "        (bn): BatchNorm1d(20, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (1): build_layer_with_layer_parameter(\n",
+       "        (padding): ConstantPad1d(padding=(33, 33), value=0)\n",
+       "        (conv1d): Conv1d(20, 400, kernel_size=(67,), stride=(1,))\n",
+       "        (bn): BatchNorm1d(400, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (2): build_layer_with_layer_parameter(\n",
+       "        (padding): ConstantPad1d(padding=(0, 1), value=0)\n",
+       "        (conv1d): Conv1d(400, 40, kernel_size=(2,), stride=(1,))\n",
+       "        (bn): BatchNorm1d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (averagepool): AdaptiveAvgPool1d(output_size=1)\n",
+       "    (hidden): Linear(in_features=40, out_features=50, bias=True)\n",
+       "  ),\n",
+       ")"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# download data from: https://drive.google.com/drive/folders/1SyX8ylC6TbwJPnLKPP_o4qdm03r-0Yby?usp=sharing\n",
+    "# dataset_path = dirname(\"../UCRArchive_2018/\")\n",
+    "# dataset_name = \"FiftyWords\"\n",
+    "\n",
+    "# load data,\n",
+    "X_train, y_train, X_test, y_test = TSC_data_loader(dataset_path, dataset_name)\n",
+    "print('train data shape', X_train.shape)\n",
+    "print()\n",
+    "print('train label shape',y_train.shape)\n",
+    "print('test data shape',X_test.shape)\n",
+    "print('test label shape',y_test.shape)\n",
+    "print('unique train label',np.unique(y_train))\n",
+    "print('unique test label',np.unique(y_test))\n",
+    "device = \"cpu\"\n",
+    "X_train, y_train, X_test, y_test = load_to_torch(X_train, y_train, X_test, y_test, device)\n",
+    "\n",
+    "# the model prints out the result every epoch\n",
+    "# defaul epoch size = 20\n",
+    "Max_kernel_size = 89\n",
+    "start_kernel_size = 1\n",
+    "# loss, optimizer, scheduler\n",
+    "input_channel = X_train.shape[1] # input channel size\n",
+    "n_class = max(y_train) + 1 # output class number\n",
+    "receptive_field_shape= min(int(X_train.shape[-1]/4),Max_kernel_size)\n",
+    "\n",
+    "model = OS_CNN_CLassifier(n_class.item(), input_channel, receptive_field_shape)\n",
+    "model.fit(X_train, y_train)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "17a15659-94d5-424c-9d64-7fb3dd2ed3b7",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.6615384615384615"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "accuracy_list=[]\n",
+    "correct = 0\n",
+    "N_test=X_test.shape[0]\n",
+    "yhat = model.predict(X_test)\n",
+    "correct += (torch.tensor(yhat) == y_test).sum().item()\n",
+    "accuracy = correct / N_test\n",
+    "accuracy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "61ff249b-d48a-4106-9409-6359c7569033",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Plot the loss and accuracy\n",
+    "# import matplotlib.pyplot as plt\n",
+    "\n",
+    "# fig, ax1 = plt.subplots()\n",
+    "# color = 'tab:red'\n",
+    "# ax1.plot(cost_list, color=color)\n",
+    "# ax1.set_xlabel('epoch', color=color)\n",
+    "# ax1.set_ylabel('Cost', color=color)\n",
+    "# ax1.tick_params(axis='y', color=color)\n",
+    "    \n",
+    "# ax2 = ax1.twinx()  \n",
+    "# color = 'tab:blue'\n",
+    "# ax2.set_ylabel('accuracy', color=color) \n",
+    "# ax2.set_xlabel('epoch', color=color)\n",
+    "# ax2.plot( accuracy_list, color=color)\n",
+    "# ax2.tick_params(axis='y', color=color)\n",
+    "# fig.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0d8d3213-c627-4580-a64d-d08eab9cb015",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/sail/models/torch/os_cnn.py b/sail/models/torch/os_cnn.py
new file mode 100644
index 0000000..a24600c
--- /dev/null
+++ b/sail/models/torch/os_cnn.py
@@ -0,0 +1,227 @@
+"""
+Omni-Scale 1D-CNN
+
+Code adapted from https://github.com/Wensi-Tang/OS-CNN
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import numpy as np
+import torch.nn as nn 
+import torch.optim as optim
+from skorch.classifier import NeuralNetClassifier
+
+
+def get_Prime_number_in_a_range(start, end):
+    Prime_list = []
+    for val in range(start, end + 1): 
+        prime_or_not = True
+        for n in range(2, val):
+            if (val % n) == 0:
+                prime_or_not = False
+                break
+        if prime_or_not:
+            Prime_list.append(val)
+    return Prime_list
+
+
+def get_out_channel_number(paramenter_layer, in_channel, prime_list):
+    out_channel_expect = int(paramenter_layer/(in_channel*sum(prime_list)))
+    return out_channel_expect
+
+
+def generate_layer_parameter_list(start,end,paramenter_number_of_layer_list, in_channel = 1):
+    prime_list = get_Prime_number_in_a_range(start, end)
+    # if prime_list == []:
+    #     # print('start = ',start, 'which is larger than end = ', end)
+    input_in_channel = in_channel
+    layer_parameter_list = []
+    for paramenter_number_of_layer in paramenter_number_of_layer_list:
+        out_channel = get_out_channel_number(paramenter_number_of_layer, in_channel, prime_list)
+        
+        tuples_in_layer= []
+        for prime in prime_list:
+            tuples_in_layer.append((in_channel,out_channel,prime))
+        in_channel =  len(prime_list)*out_channel
+        
+        layer_parameter_list.append(tuples_in_layer)
+    
+    tuples_in_layer_last = []
+    first_out_channel = len(prime_list)*get_out_channel_number(paramenter_number_of_layer_list[0], input_in_channel, prime_list)
+    tuples_in_layer_last.append((in_channel,first_out_channel,start))
+    tuples_in_layer_last.append((in_channel,first_out_channel,start+1))
+    layer_parameter_list.append(tuples_in_layer_last)
+    return layer_parameter_list
+
+
+def calculate_mask_index(kernel_length_now,largest_kernel_lenght):
+    right_zero_mast_length = math.ceil((largest_kernel_lenght-1)/2)-math.ceil((kernel_length_now-1)/2)
+    left_zero_mask_length = largest_kernel_lenght - kernel_length_now - right_zero_mast_length
+    return left_zero_mask_length, left_zero_mask_length+ kernel_length_now
+
+
+def creat_mask(number_of_input_channel,number_of_output_channel, kernel_length_now, largest_kernel_lenght):
+    ind_left, ind_right= calculate_mask_index(kernel_length_now,largest_kernel_lenght)
+    mask = np.ones((number_of_input_channel,number_of_output_channel,largest_kernel_lenght))
+    mask[:,:,0:ind_left]=0
+    mask[:,:,ind_right:]=0
+    return mask
+
+
+def creak_layer_mask(layer_parameter_list):
+    largest_kernel_lenght = layer_parameter_list[-1][-1]
+    mask_list = []
+    init_weight_list = []
+    bias_list = []
+    for i in layer_parameter_list:
+        conv = torch.nn.Conv1d(in_channels=i[0], out_channels=i[1], kernel_size=i[2])
+        ind_l,ind_r= calculate_mask_index(i[2],largest_kernel_lenght)
+        big_weight = np.zeros((i[1],i[0],largest_kernel_lenght))
+        big_weight[:,:,ind_l:ind_r]= conv.weight.detach().numpy()
+        
+        bias_list.append(conv.bias.detach().numpy())
+        init_weight_list.append(big_weight)
+        
+        mask = creat_mask(i[1],i[0],i[2], largest_kernel_lenght)
+        mask_list.append(mask)
+        
+    mask = np.concatenate(mask_list, axis=0)
+    init_weight = np.concatenate(init_weight_list, axis=0)
+    init_bias = np.concatenate(bias_list, axis=0)
+    return mask.astype(np.float32), init_weight.astype(np.float32), init_bias.astype(np.float32)
+
+    
+class build_layer_with_layer_parameter(nn.Module):
+    def __init__(self,layer_parameters):
+        super(build_layer_with_layer_parameter, self).__init__()
+
+        os_mask, init_weight, init_bias= creak_layer_mask(layer_parameters)
+        
+        
+        in_channels = os_mask.shape[1] 
+        out_channels = os_mask.shape[0] 
+        max_kernel_size = os_mask.shape[-1]
+
+        self.weight_mask = nn.Parameter(torch.from_numpy(os_mask),requires_grad=False)
+        
+        self.padding = nn.ConstantPad1d((int((max_kernel_size-1)/2), int(max_kernel_size/2)), 0)
+         
+        self.conv1d = torch.nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=max_kernel_size)
+        self.conv1d.weight = nn.Parameter(torch.from_numpy(init_weight),requires_grad=True)
+        self.conv1d.bias =  nn.Parameter(torch.from_numpy(init_bias),requires_grad=True)
+
+        self.bn = nn.BatchNorm1d(num_features=out_channels)
+    
+    def forward(self, X):
+        self.conv1d.weight.data = self.conv1d.weight*self.weight_mask
+        #self.conv1d.weight.data.mul_(self.weight_mask)
+        result_1 = self.padding(X)
+        result_2 = self.conv1d(result_1)
+        result_3 = self.bn(result_2)
+        result = F.relu(result_3)
+        return result    
+    
+    
+class _OS_CNN(nn.Module):
+    def __init__(
+        self,
+        n_class, #number of classes to predict
+        input_channel,
+        receptive_field_shape,
+        start_kernel_size = 1,
+        max_kernel_size = 89,
+        paramenter_number_of_layer_list = [8*128, 5*128*256 + 2*256*128],
+        few_shot = False
+    ):
+        super(_OS_CNN, self).__init__()
+        self.few_shot = few_shot
+        self.start_kernel_size = start_kernel_size
+        self.max_kernel_size = max_kernel_size
+        self.paramenter_number_of_layer_list = paramenter_number_of_layer_list
+        self.input_channel = input_channel
+        self.receptive_field_shape = receptive_field_shape
+        
+        layer_parameter_list = generate_layer_parameter_list(self.start_kernel_size,
+                                                     receptive_field_shape,
+                                                     self.paramenter_number_of_layer_list,
+                                                     in_channel = self.input_channel)
+        self.layer_list = []
+        
+        # building each model layer from a list of parameter
+        for i in range(len(layer_parameter_list)):
+            layer = build_layer_with_layer_parameter(layer_parameter_list[i])
+            
+            # print([layer.shape for layer in layer.parameters()])
+            self.layer_list.append(layer)
+        
+        self.net = nn.Sequential(*self.layer_list)
+            
+        self.averagepool = nn.AdaptiveAvgPool1d(1)
+        
+        out_put_channel_numebr = 0
+        for final_layer_parameters in layer_parameter_list[-1]:
+            out_put_channel_numebr = out_put_channel_numebr+ final_layer_parameters[1] 
+            
+        self.hidden = nn.Linear(out_put_channel_numebr, n_class)
+
+    def forward(self, X):
+        
+        X = self.net(X)
+
+        X = self.averagepool(X)
+        X = X.squeeze_(-1)
+
+        if not self.few_shot:
+            X = self.hidden(X)
+        return X
+    
+
+class OS_CNN_CLassifier(NeuralNetClassifier):
+    """
+    Basic OS_CNN model
+    
+    Args:
+        n_class: Number of classes to predict]
+        input_channel: Number of input channels. varies depending on data type (univariate and multivariate)
+        receptive_field_shape: Scales of receptive fields to avoid scale tuning 
+        start_kernel_size: Start kernel size. Must be a prime number 
+        max_kernel_size: Maximum kernel size. Must be a prime number
+        paramenter_number_of_layer_list: OS block layer size
+        device: cpu or gpu
+        max_epochs: Epoch number
+        batch_size: Size of the batch
+        criterion: Loss function
+        optimizer: Optimizer
+    
+    """
+    def __init__(
+        self,
+        n_class, 
+        input_channel,
+        receptive_field_shape,
+        start_kernel_size = 1,
+        max_kernel_size = 89,
+        paramenter_number_of_layer_list = [8*128, 5*128*256 + 2*256*128],
+        device = "cpu",
+        few_shot = False,
+        max_epochs = 20,
+        **kwargs):
+        super(OS_CNN_CLassifier, self).__init__(
+            module = _OS_CNN,
+            module__n_class = n_class,
+            module__input_channel = input_channel,
+            module__receptive_field_shape = receptive_field_shape,
+            module__start_kernel_size = start_kernel_size,
+            module__max_kernel_size = max_kernel_size,
+            module__paramenter_number_of_layer_list = paramenter_number_of_layer_list,
+            module__few_shot = few_shot,
+            device = device,
+            max_epochs=max_epochs, 
+            batch_size=16,
+            criterion = nn.CrossEntropyLoss,
+            optimizer = optim.Adam,
+            **kwargs
+        )
+        
\ No newline at end of file


Test patch:
diff --git a/tests/models/torch/test_os_cnn.py b/tests/models/torch/test_os_cnn.py
new file mode 100644
index 0000000..ed5c365
--- /dev/null
+++ b/tests/models/torch/test_os_cnn.py
@@ -0,0 +1,70 @@
+import pytest
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+# methods for data generation and preprocessing
+def load_to_torch(X_train, y_train, device):
+    X_train = torch.from_numpy(X_train)
+    X_train.requires_grad = False
+    X_train = X_train.to(device)
+    y_train = torch.from_numpy(y_train).to(device)
+
+    if len(X_train.shape) == 2:
+        X_train = X_train.unsqueeze_(1)
+    return X_train, y_train
+
+
+@pytest.fixture
+def Xy_train():
+    from sklearn.datasets import make_classification
+    X, y = make_classification(30, 10, random_state=0)
+    X, y = X.astype(np.float32), y.astype(np.int64)
+    
+    device = "cpu"
+    
+    X, y = load_to_torch(X, y, device)
+    return X, y
+
+
+@pytest.fixture
+def net(Xy_train):
+    from sail.models.torch.os_cnn import OS_CNN_CLassifier
+    # the model prints out the result every epoch
+    X_train, y_train = Xy_train
+    Max_kernel_size = 89
+    start_kernel_size = 1
+    input_channel = X_train.shape[1] # input channel size
+    n_class = max(y_train) + 1 # output class number
+    receptive_field_shape= min(int(X_train.shape[-1]/4),Max_kernel_size)
+    return OS_CNN_CLassifier(n_class.item(), input_channel, receptive_field_shape, max_epochs=2)
+
+
+@pytest.fixture
+def net_partial_fit(net, Xy_train):
+    X, y = Xy_train
+    for i in range(0,3):
+        net.partial_fit(X, y)
+    return net
+
+
+def test_net_learns(net, Xy_train):
+    X, y = Xy_train
+    for i in range(0,3):
+        net.partial_fit(X, y)
+    train_losses = net.history[:, 'train_loss']
+    assert train_losses[0] > train_losses[-1]
+    
+    
+def test_predict_predict_proba(net_partial_fit, Xy_train):
+    X, _ = Xy_train
+    y_pred = net_partial_fit.predict(X)
+    assert not np.allclose(y_pred, 0)
+
+    
+def test_score(net_partial_fit, Xy_train):
+    X, y = Xy_train
+    r2_score = net_partial_fit.score(X, y)
+    assert r2_score <= 1. and r2_score > 0.1
\ No newline at end of file