python - 根据数据集动态设置图层单位

Question

我尝试根据数据集更改最后一层的单位大小。这是我的代码的抽象，但它不起作用。

class cnn_model:
 num_classes = 1

 model.add(layers.Conv2D(128, (3, 3), activation='relu'))
 model.add(layers.MaxPooling2D((2, 2)))
 model.add(layers.Conv2D(256, (3, 3), activation='relu'))
 model.add(layers.MaxPooling2D((2, 2)))
 model.add(layers.Dense(num_classes, activation='softmax'))


@staticmethod
def train_two():
 cnn_mod = cnn_model
 cnn_mod.num_classes = 2
 model = cnn_mod.model

@staticmethod
def train_three():
 cnn_mod = cnn_model
 cnn_mod.num_classes = 3
 model = cnn_mod.model

标签： pythontensorflowkeras

Answer 1

当您实例化一个新类时，只需将类数作为参数传递CNNModel：

import tensorflow as tf
from tensorflow.keras import layers

class CNNModel:

    def __init__(self, num_classes=2):
        self.num_classes = num_classes
        self.model = tf.keras.models.Sequential()
        self.model.add(layers.Conv2D(128, (3, 3), activation='relu'))
        self.model.add(layers.MaxPooling2D((2, 2)))
        self.model.add(layers.Conv2D(256, (3, 3), activation='relu'))
        self.model.add(layers.MaxPooling2D((2, 2)))
        self.model.add(layers.Dense(self.num_classes, activation='softmax'))

cnnmodel = CNNModel(num_classes=3)
kerasmodel = cnnmodel.model
print(cnnmodel.num_classes) # 3   

我还建议您阅读PEP8中的命名约定和缩进。