use ignite for exp2

exp2/modified/dataset.py

@@ -3,9 +3,10 @@ import numpy
 import torch
 from torch.utils.data import DataLoader, Dataset
 from torchvision.transforms import v2 as tvtrans
+import settings
 
 class MnistDataset(Dataset):
-    """用于加载Mnist数据的自定义数据集"""
+    """适配PyTorch的自定义Dataset，用于加载MNIST数据。"""
 
     shape: int
     transform: tvtrans.Transform
@@ -29,15 +30,14 @@ class MnistDataset(Dataset):
         return self.shape
 
 
-class MnistDataSource:
-    """用于读取MNIST数据的数据读取器"""
+class MnistDataLoaders:
+    """包含适配PyTorch的训练数据Loader和测试数据Loader的类。"""
 
     train_loader: DataLoader
     test_loader: DataLoader
 
     def __init__(self, batch_size: int):
-        dataset_path = Path(__file__).resolve().parent.parent / 'datasets' / 'mnist.npz'
-        dataset = numpy.load(dataset_path)
+        dataset = numpy.load(settings.MNIST_DATASET_PATH)
 
         # 所有图片均为黑底白字
         # 6万张训练图片：60000x28x28。标签只有第一维。
@@ -50,16 +50,19 @@ class MnistDataSource:
         # 定义数据转换器
         trans = tvtrans.Compose([
             # 从uint8转换为float32并自动归一化到0-1区间
+            # YYC MARK: 下面这个被标outdated了，换下面两个替代。
             # tvtrans.ToTensor(),
             tvtrans.ToImage(),
             tvtrans.ToDtype(torch.float32, scale=True),
 
             # 为了符合后面图像的输入颜色通道条件，要在最后挤出一个新的维度
-            #tvtrans.Lambda(lambda x: x.unsqueeze(-1))
+            # YYC MARK: 上面这两步已经帮我们自动挤出那个灰度通道了。
+            # tvtrans.Lambda(lambda x: x.unsqueeze(-1))
 
             # 这个特定的标准化参数 (0.1307, 0.3081) 是 MNIST 数据集的标准化参数，这些数值是MNIST训练集的全局均值和标准差。
             # 这种标准化有助于模型训练时的数值稳定性和收敛速度。
-            #tvtrans.Normalize((0.1307,), (0.3081,)),
+            # YYC MARK: 但我不想用，反正最后训练的也收敛。
+            # tvtrans.Normalize((0.1307,), (0.3081,)),
         ])
 
         # 创建数据集
@@ -69,9 +72,9 @@ class MnistDataSource:
                                     transform=trans)
 
         # 赋值到自身
-        self.train_loader = DataLoader(dataset=train_dataset,
-                                     batch_size=batch_size,
-                                     shuffle=False)
-        self.test_loader = DataLoader(dataset=test_dataset,
-                                    batch_size=batch_size,
-                                    shuffle=False)
+        self.train_loader = DataLoader(dataset=train_dataset, 
+                                       batch_size=batch_size, 
+                                       shuffle=False)
+        self.test_loader = DataLoader(dataset=test_dataset, 
+                                      batch_size=batch_size, 
+                                      shuffle=False)

exp2/modified/model.py

@@ -21,9 +21,9 @@ class Cnn(torch.nn.Module):
         self.fc2 = torch.nn.Linear(64, 10)
 
         # 初始化模型参数
-        self._initialize_weights()
+        self.__initialize_weights()
 
-    def _initialize_weights(self):
+    def __initialize_weights(self):
         # YYC MARK:
         # 把两个全连接线性层按tensorflow默认设置初始化，即：
         # - kernel_initializer='glorot_uniform'

exp2/modified/predict.py

@@ -8,26 +8,36 @@ import matplotlib.pyplot as plt
 from model import Cnn
 
 sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
-import pytorch_gpu_utils
+import gpu_utils
 
 
 class PredictResult:
     """预测的结果"""
 
     possibilities: torch.Tensor
-    """预测结果，是每个数字不同的概率，是经过softmax后的数值"""
+    """每个数字不同的概率"""
 
     def __init__(self, possibilities: torch.Tensor):
+        """
+        创建预测结果。
+
+        :param possibilities: 传入的tensor表示每个数字不同的概率，是经过softmax后的数值。
+        其shape为二维。dim 0为batch，应当只有一维；dim 1为每个数字对应的概率。
+        """
         self.possibilities = possibilities
 
     def chosen_number(self) -> int:
         """获取最终选定的数字"""
-        # 依然是找最大的那个index
-        _, prediction = self.possibilities.max(1)
-        return prediction.item()
+        # 输出出来是10个数字各自的可能性，所以要选取最高可能性的那个对比
+        # 在dim=1上找最大的那个，就选那个。dim=0是批次所以不管他。
+        return self.possibilities.argmax(1).item()
 
     def number_possibilities(self) -> list[float]:
-        """获取每个数字出现的概率"""
+        """
+        获取每个数字出现的概率
+        
+        :return: 返回一个具有10个元素的列表，列表的每一项表示当前index所代表数字的概率。
+        """
         return list(self.possibilities[0][i].item() for i in range(10))
 
 class Predictor:
@@ -35,14 +45,14 @@ class Predictor:
     model: Cnn
 
     def __init__(self):
-        self.device = pytorch_gpu_utils.get_gpu_device()
+        self.device = gpu_utils.get_gpu_device()
         self.model = Cnn().to(self.device)
 
         # 加载保存好的模型参数
         file_path = Path(__file__).resolve().parent.parent / 'models' / 'cnn.pth'
         self.model.load_state_dict(torch.load(file_path))
 
-    def generic_predict(self, in_data: torch.Tensor) -> PredictResult:
+    def __predict(self, in_data: torch.Tensor) -> PredictResult:
         """
         其它预测函数都要使用的预测后端。其它预测函数将数据处理成Tensor，然后传递给此函数进行实际预测。
         
@@ -62,24 +72,36 @@ class Predictor:
        
 
     def predict_sketchpad(self, image: list[list[bool]]) -> PredictResult:
+        """
+        以sketchpad的数据进行预测。
+
+        :param image: 该列表的shape必须为28x28
+        """
         input = torch.Tensor(image).float()
         assert(input.dim() == 2)
         assert(input.size(0) == 28)
         assert(input.size(1) == 28)
 
-        return self.generic_predict(input)
+        return self.__predict(input)
 
     def predict_image(self, image: ImageFile.ImageFile) -> PredictResult:
-        # 确保图像为灰度图像，然后转换为numpy数组。
-        # 注意这里的numpy数组是只读的，所以要先拷贝一份
+        """
+        以Pillow图像的数据进行预测。
+        
+        :param image: Pillow图像数据。该图像必须为28x28大小。
+        """
+        # 确保图像为灰度图像，以及宽高合适
         grayscale_image = image.convert('L')
+        assert(grayscale_image.width == 28)
+        assert(grayscale_image.height == 28)
+        # 转换为numpy数组。注意这里的numpy数组是只读的，所以要先拷贝一份
         numpy_data = numpy.reshape(grayscale_image, (28, 28), copy=True)
         # 转换到Tensor，设置dtype
         data = torch.from_numpy(numpy_data).float()
         # 归一化到255，又因为图像输入是白底黑字，需要做转换。
         data.div_(255.0).sub_(1).mul_(-1)
 
-        return self.generic_predict(data)
+        return self.__predict(data)
 
 def main():
     predictor = Predictor()
@@ -101,6 +123,6 @@ def main():
 
 
 if __name__ == "__main__":
-    pytorch_gpu_utils.print_gpu_availability()
+    gpu_utils.print_gpu_availability()
     main()
     

exp2/modified/settings.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+
+MNIST_DATASET_PATH: Path = Path(__file__).resolve().parent.parent / 'datasets' / 'mnist.npz'
+"""MNIST数据集文件的路径"""
+
+SAVED_MODEL_PATH: Path = Path(__file__).resolve().parent.parent / 'models' / 'cnn.pth'
+"""训练好的模型保存的位置"""
+
+N_EPOCH: int = 5
+"""训练时的epoch次数"""
+N_BATCH_SIZE: int = 1000
+"""训练时的batch size"""

exp2/modified/sketchpad.py

@@ -5,7 +5,7 @@ import tkinter as tk
 from predict import PredictResult, Predictor
 
 sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
-import pytorch_gpu_utils
+import gpu_utils
 
 
 class SketchpadApp:
@@ -181,7 +181,7 @@ class SketchpadApp:
     # endregion
 
 if __name__ == "__main__":
-    pytorch_gpu_utils.print_gpu_availability()
+    gpu_utils.print_gpu_availability()
     
     predictor = Predictor()
 

exp2/modified/train.py

@@ -6,188 +6,184 @@ import torchinfo
 import ignite.engine
 import ignite.metrics
 from ignite.engine import Engine, Events
+from ignite.metrics import Accuracy, Loss
 from ignite.handlers.tqdm_logger import ProgressBar
-from dataset import MnistDataSource
+from dataset import MnistDataLoaders
 from model import Cnn
+import settings
 
 sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
-import pytorch_gpu_utils
+import gpu_utils
 
 
-class Trainer:
-    N_EPOCH: typing.ClassVar[int] = 5
-    N_BATCH_SIZE: typing.ClassVar[int] = 1000
-
-    device: torch.device
-    data_source: MnistDataSource
-    model: Cnn
-
-    def __init__(self):
-        self.device = pytorch_gpu_utils.get_gpu_device()
-        self.data_source = MnistDataSource(Trainer.N_BATCH_SIZE)
-        self.model = Cnn().to(self.device)
-        # 展示模型结构。批次为指定批次数量，通道只有一个灰度通道，大小28x28。
-        torchinfo.summary(self.model, (Trainer.N_BATCH_SIZE, 1, 28, 28))
-
-    def train(self):
-        optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
-        # optimizer = torch.optim.AdamW(
-        #     self.model.parameters(),
-        #     lr=0.001,          # 两者默认学习率都是 0.001
-        #     betas=(0.9, 0.999), # 两者默认值相同
-        #     eps=1e-07,         # 【关键】匹配 TensorFlow 的默认 epsilon
-        #     weight_decay=0.0,   # 两者默认都是 0
-        #     amsgrad=False       # 两者默认都是 False
-        # )
-        loss_func = torch.nn.CrossEntropyLoss()
-
-        for epoch in range(Trainer.N_EPOCH):
-            self.model.train()
-
-            batch_images: torch.Tensor
-            batch_labels: torch.Tensor
-            for batch_index, (batch_images, batch_labels) in enumerate(self.data_source.train_loader):
-                gpu_images = batch_images.to(self.device)
-                gpu_labels = batch_labels.to(self.device)
-
-                prediction: torch.Tensor = self.model(gpu_images)
-                loss: torch.Tensor = loss_func(prediction, gpu_labels)
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
-
-                if batch_index % 100 == 0:
-                    literal_loss = loss.item()
-                    print(f'Epoch: {epoch+1}, Batch: {batch_index}, Loss: {literal_loss:.4f}')
-
-    def save(self):
-        file_dir_path = Path(__file__).resolve().parent.parent / 'models'
-        file_dir_path.mkdir(parents=True, exist_ok=True)
-        file_path = file_dir_path / 'cnn.pth'
-        torch.save(self.model.state_dict(), file_path)
-        print(f'模型已保存至：{file_path}')
-
-    def test(self):
-        self.model.eval()
-        correct_sum = 0
-        total_sum = 0
-
-        with torch.no_grad():
-            for batch_images, batch_labels in self.data_source.test_loader:
-                gpu_images = batch_images.to(self.device)
-                gpu_labels = batch_labels.to(self.device)
-
-                possibilities: torch.Tensor = self.model(gpu_images)
-                # 输出出来是10个数字各自的可能性，所以要选取最高可能性的那个对比
-                # 在dim=1上找最大的那个，就选那个。dim=0是批次所以不管他。
-                _, prediction = possibilities.max(1)
-                # 返回标签的个数作为这一批的总个数
-                total_sum += gpu_labels.size(0)
-                correct_sum += prediction.eq(gpu_labels).sum()
-
-        test_acc = 100. * correct_sum / total_sum
-        print(f"准确率: {test_acc:.4f}%，共测试了{total_sum}张图片")
-
-def main():
-    trainer = Trainer()
-    trainer.train()
-    trainer.save()
-    trainer.test()
-
-# N_EPOCH: int = 5
-# N_BATCH_SIZE: int = 1000
-# N_LOG_INTERVAL: int = 10
-
 # class Trainer:
 
 #     device: torch.device
-#     data_source: MnistDataSource
+#     data_source: MnistDataLoaders
 #     model: Cnn
 
-#     trainer: Engine
-#     train_evaluator: Engine
-#     test_evaluator: Engine
-
 #     def __init__(self):
 #         self.device = gpu_utils.get_gpu_device()
+#         self.data_source = MnistDataLoaders(Trainer.N_BATCH_SIZE)
 #         self.model = Cnn().to(self.device)
-#         self.data_source = MnistDataSource(batch_size=N_BATCH_SIZE)
 #         # 展示模型结构。批次为指定批次数量，通道只有一个灰度通道，大小28x28。
-#         torchinfo.summary(self.model, (N_BATCH_SIZE, 1, 28, 28))
+#         torchinfo.summary(self.model, (Trainer.N_BATCH_SIZE, 1, 28, 28))
 
-#         #optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
-#         optimizer = torch.optim.AdamW(
-#             self.model.parameters(),
-#             lr=0.001,          # 两者默认学习率都是 0.001
-#             betas=(0.9, 0.999), # 两者默认值相同
-#             eps=1e-07,         # 【关键】匹配 TensorFlow 的默认 epsilon
-#             weight_decay=0.0,   # 两者默认都是 0
-#             amsgrad=False       # 两者默认都是 False
-#         )
-#         criterion = torch.nn.CrossEntropyLoss()
-#         self.trainer = ignite.engine.create_supervised_trainer(
-#             self.model, optimizer, criterion, self.device
-#         )
+#     def train(self):
+#         optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
+#         # optimizer = torch.optim.AdamW(
+#         #     self.model.parameters(),
+#         #     lr=0.001,          # 两者默认学习率都是 0.001
+#         #     betas=(0.9, 0.999), # 两者默认值相同
+#         #     eps=1e-07,         # 【关键】匹配 TensorFlow 的默认 epsilon
+#         #     weight_decay=0.0,   # 两者默认都是 0
+#         #     amsgrad=False       # 两者默认都是 False
+#         # )
+#         loss_func = torch.nn.CrossEntropyLoss()
 
-#         eval_metrics = {
-#             "accuracy": ignite.metrics.Accuracy(device=self.device),
-#             "loss": ignite.metrics.Loss(criterion, device=self.device)
-#         }
-#         self.train_evaluator = ignite.engine.create_supervised_evaluator(
-#             self.model, metrics=eval_metrics, device=self.device)
-#         self.test_evaluator = ignite.engine.create_supervised_evaluator(
-#             self.model, metrics=eval_metrics, device=self.device)
-        
-#         self.trainer.add_event_handler(
-#             Events.ITERATION_COMPLETED(every=N_LOG_INTERVAL),
-#             lambda engine: self.log_intrain_loss(engine)
-#         )
-#         self.trainer.add_event_handler(
-#             Events.EPOCH_COMPLETED,
-#             lambda trainer: self.log_train_results(trainer)
-#         )
-#         self.trainer.add_event_handler(
-#             Events.COMPLETED,
-#             lambda _: self.log_test_results()
-#         )
-#         self.trainer.add_event_handler(
-#             Events.COMPLETED,
-#             lambda _: self.save_model()
-#         )
+#         for epoch in range(Trainer.N_EPOCH):
+#             self.model.train()
 
-#         progressbar = ProgressBar()
-#         progressbar.attach(self.trainer)
+#             batch_images: torch.Tensor
+#             batch_labels: torch.Tensor
+#             for batch_index, (batch_images, batch_labels) in enumerate(self.data_source.train_loader):
+#                 gpu_images = batch_images.to(self.device)
+#                 gpu_labels = batch_labels.to(self.device)
 
-#     def run(self):
-#         self.trainer.run(self.data_source.train_loader, max_epochs=N_EPOCH)
+#                 prediction: torch.Tensor = self.model(gpu_images)
+#                 loss: torch.Tensor = loss_func(prediction, gpu_labels)
+#                 optimizer.zero_grad()
+#                 loss.backward()
+#                 optimizer.step()
 
-#     def log_intrain_loss(self, engine: Engine):
-#         print(f"Epoch: {engine.state.epoch}, Loss: {engine.state.output:.4f}\r", end="")
+#                 if batch_index % 100 == 0:
+#                     literal_loss = loss.item()
+#                     print(f'Epoch: {epoch+1}, Batch: {batch_index}, Loss: {literal_loss:.4f}')
 
-#     def log_train_results(self, trainer: Engine):
-#         self.train_evaluator.run(self.data_source.train_loader)
-#         metrics = self.train_evaluator.state.metrics
-#         print()
-#         print(f"Training - Epoch: {trainer.state.epoch}, Avg Accuracy: {metrics['accuracy']:.4f}, Avg Loss: {metrics['loss']:.4f}")
-
-#     def log_test_results(self):
-#         self.test_evaluator.run(self.data_source.test_loader)
-#         metrics = self.test_evaluator.state.metrics
-#         print(f"Test - Avg Accuracy: {metrics['accuracy']:.4f} Avg Loss: {metrics['loss']:.4f}")
-
-#     def save_model(self):
+#     def save(self):
 #         file_dir_path = Path(__file__).resolve().parent.parent / 'models'
 #         file_dir_path.mkdir(parents=True, exist_ok=True)
 #         file_path = file_dir_path / 'cnn.pth'
 #         torch.save(self.model.state_dict(), file_path)
-#         print(f'Model was saved into: {file_path}')
+#         print(f'模型已保存至：{file_path}')
+
+#     def test(self):
+#         self.model.eval()
+#         correct_sum = 0
+#         total_sum = 0
+
+#         with torch.no_grad():
+#             batch_images: torch.Tensor
+#             batch_labels: torch.Tensor
+#             for batch_images, batch_labels in self.data_source.test_loader:
+#                 gpu_images = batch_images.to(self.device)
+#                 gpu_labels = batch_labels.to(self.device)
+
+#                 possibilities: torch.Tensor = self.model(gpu_images)
+#                 # 输出出来是10个数字各自的可能性，所以要选取最高可能性的那个对比
+#                 # 在dim=1上找最大的那个，就选那个。dim=0是批次所以不管他。
+#                 _, prediction = possibilities.max(1)
+#                 # 返回标签的个数作为这一批的总个数
+#                 total_sum += gpu_labels.size(0)
+#                 correct_sum += prediction.eq(gpu_labels).sum()
+
+#         test_acc = 100. * correct_sum / total_sum
+#         print(f"准确率: {test_acc:.4f}%，共测试了{total_sum}张图片")
 
 # def main():
 #     trainer = Trainer()
-#     trainer.run()
+#     trainer.train()
+#     trainer.save()
+#     trainer.test()
 
+class Trainer:
+    """核心训练器"""
+
+    device: torch.device
+    data_source: MnistDataLoaders
+    model: Cnn
+
+    trainer: Engine
+    trainer_accuracy: Accuracy
+    evaluator: Engine
+    pbar: ProgressBar
+
+    def __init__(self):
+        # 创建训练设备，模型和数据加载器。
+        self.device = gpu_utils.get_gpu_device()
+        self.model = Cnn().to(self.device)
+        self.data_source = MnistDataLoaders(batch_size=settings.N_BATCH_SIZE)
+        # 展示模型结构。批次为指定批次数量，通道只有一个灰度通道，大小28x28。
+        torchinfo.summary(self.model, (settings.N_BATCH_SIZE, 1, 28, 28))
+
+        # 优化器和损失函数
+        optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
+        criterion = torch.nn.CrossEntropyLoss()
+        # 创建训练器
+        self.trainer = ignite.engine.create_supervised_trainer(
+            self.model, optimizer, criterion, self.device,
+            # 输出转换为这种形式，因为后面的测量需要用其中一些参数。
+            # 默认的输出只输出loss。
+            output_transform=lambda x, y, y_pred, loss: {
+                'loss': loss.item(),
+                'y': y,
+                'y_pred': y_pred
+            }
+        )
+        # 设置训练器测量数据
+        self.trainer_accuracy = Accuracy(
+            device=self.device,
+            # 转换为Accuracy需要的形式。
+            output_transform=lambda o: (o['y_pred'], o['y'])
+        )
+        self.trainer_accuracy.attach(self.trainer, 'accuracy')
+        # YYC MARK: 这里要手动reset一下，不然第一次运行没有accuracy
+        self.trainer_accuracy.reset()
+        # 每次epoch前重置accuracy
+        self.trainer.add_event_handler(
+            Events.EPOCH_STARTED,
+            lambda: self.trainer_accuracy.reset()
+        )
+        # 将训练器关联到进度条
+        self.pbar = ProgressBar(persist=True)
+        self.pbar.attach(self.trainer, 
+                         metric_names=['accuracy'],
+                         output_transform=lambda o: {"loss": o['loss']})
+        # 训练完毕后保存模型
+        self.trainer.add_event_handler(
+            Events.COMPLETED,
+            lambda: self.save_model()
+        )
+
+        # 创建测试的评估器的评估量
+        evaluator_metrics = {
+            "accuracy": ignite.metrics.Accuracy(device=self.device),
+            "loss": ignite.metrics.Loss(criterion, device=self.device)
+        }
+        # 创建测试评估器
+        self.evaluator = ignite.engine.create_supervised_evaluator(
+            self.model, metrics=evaluator_metrics, device=self.device)
+        
+    def run(self):
+        # 训练模型
+        self.trainer.run(self.data_source.train_loader, max_epochs=settings.N_EPOCH)
+        # 测试模型并输出结果
+        self.evaluator.run(self.data_source.test_loader)
+        metrics = self.evaluator.state.metrics
+        print(f"Test Done. Accuracy: {metrics['accuracy']:.4f} Loss: {metrics['loss']:.4f}")
+
+    def save_model(self):
+        # 确保保存模型的文件夹存在。
+        settings.SAVED_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
+        # 仅保存模型参数
+        torch.save(self.model.state_dict(), settings.SAVED_MODEL_PATH)
+        print(f'Model was saved into: {settings.SAVED_MODEL_PATH}')
+
+def main():
+    trainer = Trainer()
+    trainer.run()
 
 
 if __name__ == "__main__":
-    pytorch_gpu_utils.print_gpu_availability()
+    gpu_utils.print_gpu_availability()
     main()

exp2/source/predict.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 import tensorflow as tf
 from PIL import Image
 import numpy as np
@@ -33,6 +34,7 @@ class Predict(object):
 
 if __name__ == "__main__":
     app = Predict()
-    app.predict('./test_images/0.png') 
-    app.predict('./test_images/1.png')
-    app.predict('./test_images/4.png')
+    images_dir = Path(__file__).resolve().parent.parent / 'test_images'
+    app.predict(images_dir / '0.png') 
+    app.predict(images_dir / '1.png')
+    app.predict(images_dir / '4.png')

exp2/source/train.py

@@ -1,12 +1,6 @@
 from pathlib import Path
-import sys
 import tensorflow as tf
-import keras
-from keras import datasets, layers, models
-
-sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
-import tensorflow_gpu_util
-
+from tensor.keras import datasets, layers, models
 
 class CNN(object):
     def __init__(self):
@@ -46,7 +40,7 @@ class Train:
     def train(self):
         check_path = Path(__file__).resolve().parent.parent / 'models' / 'cnn.ckpt'
         # period 每隔5epoch保存一次
-        save_model_cb = keras.callbacks.ModelCheckpoint(
+        save_model_cb = tf.keras.callbacks.ModelCheckpoint(
             str(check_path), save_weights_only=True, verbose=1, period=5)
         self.cnn.model.compile(optimizer='adam',
                                loss='sparse_categorical_crossentropy',
@@ -58,6 +52,5 @@ class Train:
         print("准确率: %.4f, 共测试了%d张图片 " % (test_acc, len(self.data.test_labels)))
 
 if __name__ == "__main__":
-    tensorflow_gpu_util.print_gpu_availability()
-    #app = Train()
-    #app.train()
+    app = Train()
+    app.train()