demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 轴承故障诊断结果演示"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 导入必要的外部库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader, random_split,Dataset\n",
    "from scipy import io\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义CNN网络结构模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    '''定义一维卷积神经网络模型'''\n",
    "\n",
    "    def __init__(self, in_channel=3, out_channel=9):\n",
    "        super(CNN, self).__init__()\n",
    "        '''除输入层外，每个层级都包含了卷积、激活和池化三层'''\n",
    "        '''输出层额外包含了BatchNorm层，提高网络收敛速度以及稳定性'''\n",
    "        '''第一层卷积核大小为64，之后逐层递减'''\n",
    "        self.layer1 = nn.Sequential(\n",
    "            nn.Conv1d(in_channel, 16, kernel_size=64, stride=16, padding=24),\n",
    "            nn.BatchNorm1d(16),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.layer2 = nn.Sequential(\n",
    "            nn.Conv1d(16, 32, kernel_size=16, padding=8),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2))\n",
    "\n",
    "        self.layer3 = nn.Sequential(\n",
    "            nn.Conv1d(32, 48, kernel_size=8, padding=4),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "        self.layer4 = nn.Sequential(\n",
    "            nn.Conv1d(48, 64, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.layer5 = nn.Sequential(\n",
    "            nn.Conv1d(64, 64, kernel_size=3),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "        # 全连接层定义 引入Dropout机制以提高泛化能力\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(64, 128),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(128, out_channel)\n",
    "        )\n",
    "        # 使用softmax函数以计算输出从属于每一类的概率\n",
    "        self.softmax = nn.Softmax(1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        '''前向传播'''\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "        x = self.layer5(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc(x)\n",
    "        x = x.view(x.size(0), x.size(1), 1)\n",
    "        x = self.softmax(x)\n",
    "        return x\n",
    "\n",
    "class LSTM_CNN(nn.Module):\n",
    "    '''定义LSTM-CNN网络模型'''\n",
    "\n",
    "    def __init__(self, DEVICE, in_channel=3, out_channel=9):\n",
    "        super(LSTM_CNN, self).__init__()\n",
    "        self.DEVICE = DEVICE\n",
    "        '''LSTM相关神经元定义'''\n",
    "        self.lstm_layer1 = nn.LSTM(in_channel, 32)\n",
    "        self.lstm_layer2 = nn.LSTM(64, 1)\n",
    "        self.lstm_fc1 = nn.Sequential(\n",
    "            nn.Linear(32, 64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(64, 64)\n",
    "        )\n",
    "\n",
    "        self.lstm_fc2 = nn.Sequential(\n",
    "            nn.Linear(1024, 256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(256, 64)\n",
    "        )\n",
    "\n",
    "        '''CNN相关神经元定义'''\n",
    "        self.cnn_layer1 = nn.Sequential(\n",
    "            nn.Conv1d(in_channel, 16, kernel_size=64, stride=16, padding=24),\n",
    "            nn.BatchNorm1d(16),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.cnn_layer2 = nn.Sequential(\n",
    "            nn.Conv1d(16, 32, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm1d(32),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2))\n",
    "\n",
    "        self.cnn_layer3 = nn.Sequential(\n",
    "            nn.Conv1d(32, 64, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.cnn_layer4 = nn.Sequential(\n",
    "            nn.Conv1d(64, 64, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.cnn_layer5 = nn.Sequential(\n",
    "            nn.Conv1d(64, 64, kernel_size=3),\n",
    "            nn.BatchNorm1d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=2)\n",
    "        )\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(128, 64),\n",
    "            nn.Dropout(0.3),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(64, out_channel)\n",
    "        )\n",
    "        # 使用softmax函数以计算输出从属于每一类的概率\n",
    "        self.softmax = nn.Softmax(1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        '''前向传播'''\n",
    "        '''*******LSTM*******'''\n",
    "        # 初始化隐藏神经元\n",
    "        x_lstm = x.permute(0, 2, 1)\n",
    "        x_lstm.to(self.DEVICE)\n",
    "        h1 = torch.zeros(1, 128, 32).to(self.DEVICE)\n",
    "        c1 = torch.zeros_like(h1).to(self.DEVICE)\n",
    "        h2 = torch.zeros(1, 128, 1).to(self.DEVICE)\n",
    "        c2 = torch.zeros_like(h2).to(self.DEVICE)\n",
    "        y_lstm_ = []\n",
    "        # 对原时序信号分段\n",
    "        for i in range(8):\n",
    "            x_lstm_ = x_lstm[:, i*128:(i+1)*128]\n",
    "            y, (h1, c1) = self.lstm_layer1(x_lstm_, (h1, c1))\n",
    "            y = self.lstm_fc1(y)\n",
    "            y, (h2, c2) = self.lstm_layer2(y, (h2, c2))\n",
    "            y.to(self.DEVICE)\n",
    "            y_lstm_.append(y)\n",
    "        # 合并每一段的结果\n",
    "        y_lstm = torch.cat(y_lstm_, 1)\n",
    "        y_lstm = y_lstm.view(y_lstm.size(0), -1)\n",
    "        y_lstm = self.lstm_fc2(y_lstm)\n",
    "        '''*******CNN*******'''\n",
    "        x = self.cnn_layer1(x)\n",
    "        x = self.cnn_layer2(x)\n",
    "        x = self.cnn_layer3(x)\n",
    "        x = self.cnn_layer4(x)\n",
    "        x = self.cnn_layer5(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        '''******LSTM+CNN******'''\n",
    "        # 连接LSTM和CNN的输出，并通过全连接神经元\n",
    "        x = torch.cat([x, y_lstm], 1)\n",
    "        x = self.fc(x)\n",
    "        x = x.view(x.size(0), x.size(1), 1)\n",
    "        y = self.softmax(x)\n",
    "        return y\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义数据集类以及必要的工具函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Data_set(Dataset):\n",
    "    def __init__(self, data, label):\n",
    "        super().__init__()\n",
    "        self.x = torch.tensor(data)\n",
    "        self.y = torch.tensor(label).view(-1, 1)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        assert idx < len(self.y)\n",
    "        return self.x[idx], self.y[idx]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "\n",
    "\n",
    "def file_read(filename: str):\n",
    "    '''读取数据'''\n",
    "    label = filename.split('.')[0].split('/')[-1].split(\"\\\\\")[-1]\n",
    "    data = pd.read_csv(filename)\n",
    "    array = data.to_numpy(dtype=np.float32)\n",
    "    return get_idx(label), array\n",
    "\n",
    "\n",
    "def get_idx(label: str):\n",
    "    '''根据标签获得其索引'''\n",
    "    labels = ['BL07', 'BL14', 'BL21', 'IR07',\n",
    "              'IR14', 'IR21', 'OR07', 'OR14', 'OR21']\n",
    "    idx = labels.index(label)\n",
    "    return idx\n",
    "\n",
    "\n",
    "def get_label(idx):\n",
    "    '''根据索引获得标签内容'''\n",
    "    labels = ['BL07', 'BL14', 'BL21', 'IR07',\n",
    "              'IR14', 'IR21', 'OR07', 'OR14', 'OR21']\n",
    "    return labels[idx]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 获取数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DEVICE:cuda\n"
     ]
    }
   ],
   "source": [
    "global DEVICE, BATCH_SIZE, EPOCHS, LENGTH, PATH, LOSSES, ACCS\n",
    "LENGTH = 1024\n",
    "BATCH_SIZE = 64\n",
    "EPOCHS = 20\n",
    "STRIDE = 128\n",
    "PATH = os.path.abspath('./../data')\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "LOSSES = []\n",
    "ACCS = []\n",
    "print(f\"DEVICE:{DEVICE}\")\n",
    "files = os.listdir(PATH)\n",
    "label = []\n",
    "data = []\n",
    "for file in files:\n",
    "    if \"csv\" not in file:\n",
    "        continue\n",
    "    labeli, xi = file_read(PATH + \"/\" + file)\n",
    "    for j in range(0, len(xi)-LENGTH, STRIDE):\n",
    "        label.append(labeli)\n",
    "        data.append(xi[j:j+LENGTH, :].T)\n",
    "    label.append(labeli)\n",
    "    data.append(xi[-LENGTH:, :].T)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = Data_set(data, label)\n",
    "'''数据集分割'''\n",
    "train_size = int(0.7*len(ds))\n",
    "test_size = len(ds) - train_size\n",
    "train_loader, test_loader = random_split(ds, [train_size, test_size])\n",
    "train_loader = DataLoader(train_loader, BATCH_SIZE, shuffle=True)\n",
    "test_loader = DataLoader(test_loader, BATCH_SIZE, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义训练及验证函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(net: CNN, dataloader: DataLoader, loss_func, optimizer: optim.Adam, epoch: int):\n",
    "    global DEVICE\n",
    "    net.train()\n",
    "    for i, (x, label) in enumerate(dataloader):\n",
    "        x, label = x.to(DEVICE), label.to(DEVICE)\n",
    "        optimizer.zero_grad()\n",
    "        y = net(x).to(DEVICE)\n",
    "        loss = loss_func(y, label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if ((epoch+1) % 10 == 0) and ((i+1) % 8 == 0):\n",
    "            print(\"Train Epoch: {}\\t[{: >4}/{} ({:.0f}%)] Loss: {:.6f}\".format(\n",
    "                epoch+1, i*len(x), len(dataloader.dataset), 100*i/len(dataloader), loss.item()))\n",
    "\n",
    "\n",
    "def test(net: CNN, dataloader: DataLoader, loss_func, datatype: str):\n",
    "    global DEVICE\n",
    "    net.eval()\n",
    "    test_loss = 0\n",
    "    cnt = 0\n",
    "    with torch.no_grad():\n",
    "        for x, label in dataloader:\n",
    "            x, label = x.to(DEVICE), label.to(DEVICE)\n",
    "            y = net(x).to(DEVICE)\n",
    "            test_loss += loss_func(y, label)\n",
    "            predict = y.max(1, keepdim=True)[1]\n",
    "            cnt += predict.eq(label.view_as(predict)).sum().item()\n",
    "    test_loss /= len(dataloader.dataset)\n",
    "    print(\"{} Data:  Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\".format(datatype,\n",
    "                                                                             test_loss, cnt, len(dataloader.dataset), 100*cnt/len(dataloader.dataset)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 训练卷积神经网络"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "训练一维卷积网络，输出其loss以及其在数据集上的准确率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 10\t[ 448/4166 (11%)] Loss: 1.389482\n",
      "Train Epoch: 10\t[ 960/4166 (23%)] Loss: 1.372049\n",
      "Train Epoch: 10\t[1472/4166 (35%)] Loss: 1.372186\n",
      "Train Epoch: 10\t[1984/4166 (47%)] Loss: 1.371996\n",
      "Train Epoch: 10\t[2496/4166 (59%)] Loss: 1.372764\n",
      "Train Epoch: 10\t[3008/4166 (71%)] Loss: 1.385443\n",
      "Train Epoch: 10\t[3520/4166 (83%)] Loss: 1.372197\n",
      "Train Epoch: 10\t[4032/4166 (95%)] Loss: 1.387763\n",
      "\n",
      "训练 Data:  Average loss: 0.0218, Accuracy: 4157/4166 (100%)\n",
      "验证 Data:  Average loss: 0.0216, Accuracy: 1779/1786 (100%)\n",
      "\n",
      "Train Epoch: 20\t[ 448/4166 (11%)] Loss: 1.371956\n",
      "Train Epoch: 20\t[ 960/4166 (23%)] Loss: 1.401479\n",
      "Train Epoch: 20\t[1472/4166 (35%)] Loss: 1.372135\n",
      "Train Epoch: 20\t[1984/4166 (47%)] Loss: 1.404388\n",
      "Train Epoch: 20\t[2496/4166 (59%)] Loss: 1.371960\n",
      "Train Epoch: 20\t[3008/4166 (71%)] Loss: 1.372253\n",
      "Train Epoch: 20\t[3520/4166 (83%)] Loss: 1.372465\n",
      "Train Epoch: 20\t[4032/4166 (95%)] Loss: 1.473830\n",
      "\n",
      "训练 Data:  Average loss: 0.0220, Accuracy: 4094/4166 (98%)\n",
      "验证 Data:  Average loss: 0.0219, Accuracy: 1742/1786 (98%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "net = CNN(3, 9).to(DEVICE)\n",
    "optimizer = optim.Adam(net.parameters())\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    train(net, train_loader, loss_func, optimizer, epoch)\n",
    "    if (epoch+1) % 10 == 0:\n",
    "        print()\n",
    "        test(net,  train_loader, loss_func, \"训练\")\n",
    "        test(net, test_loader, loss_func, \"验证\")\n",
    "        print()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "训练结合LSTM的一维卷积神经网络，输出其LOSS以及在训练集和验证集上的准确率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 10\t[ 448/4166 (11%)] Loss: 1.382450\n",
      "Train Epoch: 10\t[ 960/4166 (23%)] Loss: 1.385502\n",
      "Train Epoch: 10\t[1472/4166 (35%)] Loss: 1.385053\n",
      "Train Epoch: 10\t[1984/4166 (47%)] Loss: 1.376373\n",
      "Train Epoch: 10\t[2496/4166 (59%)] Loss: 1.374391\n",
      "Train Epoch: 10\t[3008/4166 (71%)] Loss: 1.374265\n",
      "Train Epoch: 10\t[3520/4166 (83%)] Loss: 1.373785\n",
      "Train Epoch: 10\t[4032/4166 (95%)] Loss: 1.373246\n",
      "\n",
      "训练 Data:  Average loss: 0.0217, Accuracy: 4166/4166 (100%)\n",
      "验证 Data:  Average loss: 0.0215, Accuracy: 1786/1786 (100%)\n",
      "\n",
      "Train Epoch: 20\t[ 448/4166 (11%)] Loss: 1.372115\n",
      "Train Epoch: 20\t[ 960/4166 (23%)] Loss: 1.372114\n",
      "Train Epoch: 20\t[1472/4166 (35%)] Loss: 1.372649\n",
      "Train Epoch: 20\t[1984/4166 (47%)] Loss: 1.372187\n",
      "Train Epoch: 20\t[2496/4166 (59%)] Loss: 1.376122\n",
      "Train Epoch: 20\t[3008/4166 (71%)] Loss: 1.372383\n",
      "Train Epoch: 20\t[3520/4166 (83%)] Loss: 1.372622\n",
      "Train Epoch: 20\t[4032/4166 (95%)] Loss: 1.372451\n",
      "\n",
      "训练 Data:  Average loss: 0.0217, Accuracy: 4166/4166 (100%)\n",
      "验证 Data:  Average loss: 0.0215, Accuracy: 1786/1786 (100%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "net = LSTM_CNN(DEVICE).to(DEVICE)\n",
    "optimizer = optim.Adam(net.parameters())\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    train(net, train_loader, loss_func, optimizer, epoch)\n",
    "    if (epoch+1) % 10 == 0:\n",
    "        print()\n",
    "        test(net,  train_loader, loss_func, \"训练\")\n",
    "        test(net, test_loader, loss_func, \"验证\")\n",
    "        print()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成测试集标签"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>index</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>BL21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>IR14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>OR14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>OR14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>OR07</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1795</th>\n",
       "      <td>1795</td>\n",
       "      <td>IR14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1796</th>\n",
       "      <td>1796</td>\n",
       "      <td>IR21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1797</th>\n",
       "      <td>1797</td>\n",
       "      <td>IR14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1798</th>\n",
       "      <td>1798</td>\n",
       "      <td>IR21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1799</th>\n",
       "      <td>1799</td>\n",
       "      <td>IR21</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1800 rows × 2 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      index label\n",
       "0         0  BL21\n",
       "1         1  IR14\n",
       "2         2  OR14\n",
       "3         3  OR14\n",
       "4         4  OR07\n",
       "...     ...   ...\n",
       "1795   1795  IR14\n",
       "1796   1796  IR21\n",
       "1797   1797  IR14\n",
       "1798   1798  IR21\n",
       "1799   1799  IR21\n",
       "\n",
       "[1800 rows x 2 columns]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = io.loadmat(PATH+\"\\\\Fault_Diag_Data.mat\")\n",
    "test_data = torch.tensor(\n",
    "    data[\"TestDataArray\"], dtype=torch.float32).permute(2, 1, 0)\n",
    "net.eval()\n",
    "predict = []\n",
    "for x in test_data:\n",
    "    x = x.reshape(1, x.shape[0], x.shape[1]).to(DEVICE)\n",
    "    y = net(x)\n",
    "    output = y.max(1, keepdim=True)[1]\n",
    "    predict.append(get_label(output.item()))\n",
    "\n",
    "df = {\n",
    "    \"index\": [i for i in range(len(test_data))],\n",
    "    \"label\": predict\n",
    "}\n",
    "\n",
    "df = pd.DataFrame(df)\n",
    "# df.to_excel(\"results.xlsx\", index=False)\n",
    "df\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.1 64-bit",
   "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.1"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "11938c6bc6919ae2720b4d5011047913343b08a43b18698fd82dedb0d4417594"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}