减少CNN训练时，测试集的内存占用

2019-05-03 10:50:10 +08:00
parent 86e04d444c
commit 174809e6a7
3 changed files with 18 additions and 13 deletions
--- a/tools/TrainCNN/generate.py
+++ b/tools/TrainCNN/generate.py
@@ -67,6 +67,15 @@ class DataSet:
            labels.append(self.train_labels[id])
        return np.array(samples), np.array(labels)

+    def sample_test_sets(self, length):
+        samples = []
+        labels = []
+        for i in range(length):
+            id = random.randint(0, len(self.test_samples)-1)
+            samples.append(self.test_samples[id])
+            labels.append(self.test_labels[id])
+        return np.array(samples), np.array(labels)
+
    def all_train_sets(self):
        return self.train_samples[:], self.train_labels[:]