fix test_bench.py

Author: dunzeng

tests/test_core/task_setting_for_test.py

@@ -4,6 +4,28 @@
 from torch.utils.data.dataset import Dataset
 from fedlab.core.client.trainer import ClientTrainer
 
+
+class CNN_Mnist(nn.Module):
+    def __init__(self):
+        super(CNN_Mnist, self).__init__()
+        self.conv1 = nn.Conv2d(1, 32, kernel_size=(5, 5))
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=(5, 5))
+        self.pool = nn.MaxPool2d(kernel_size=(2, 2))
+        self.fc1 = nn.Linear(in_features=1024, out_features=512)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(512, 10)
+
+    def forward(self, x):
+        x = self.pool(self.conv1(x))
+        x = self.pool(self.conv2(x))
+        x = x.view(x.shape[0], -1)
+        x = self.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+
+
 model = nn.Sequential(nn.Linear(20,10), nn.Linear(10,5))
 optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
 criterion = nn.L1Loss()

tests/test_core/test_server/test_handler.py

@@ -17,7 +17,7 @@
 
 import torch
 
-from fedlab_benchmarks.models.cnn import CNN_Mnist
+from ..task_setting_for_test import CNN_Mnist
 from fedlab.utils.serialization import SerializationTool
 from fedlab.core.server.handler import AsyncParameterServerHandler, SyncParameterServerHandler
 

tests/test_utils/test_compressor.py

@@ -14,7 +14,7 @@
 
 import unittest
 
-from fedlab_benchmarks.models.cnn import CNN_Mnist
+from ..test_core.task_setting_for_test import CNN_Mnist
 from fedlab.utils.compressor import TopkCompressor