使用在imagnet上训练的模型，如果原始imagenet数据集中包含图像，则可以很好地识别图像。如果图像不作为类出现，则模型的性能将非常差。通常，您需要为数据集中的唯一类自定义模型。此过程称为迁移学习。首先，你必须决定你想拥有什么类，并收集与每个类相关联的适当图像。例如，假设你有警车、校车、消防车垃圾车和送货货车。所以你需要为每个类收集适当的图像。通常你至少需要为每个类收集大约120到150个图像。所以我们现在有5个类。创建一个单独的目录调用是sdir。在sdir下为每个类创建5个子目录。将它们命名为警车、校车、现在可以使用下面的函数将数据集拆分为train_df、test_df和valid_df三个数据集。

def preprocess (sdir, trsplit, vsplit):
    filepaths=[]
    labels=[]    
    classlist=os.listdir(sdir)
    for klass in classlist:
        classpath=os.path.join(sdir,klass)
        if os.path.isdir(classpath):
            flist=os.listdir(classpath)
            for f in flist:
                fpath=os.path.join(classpath,f)
                filepaths.append(fpath)
                labels.append(klass)
    Fseries=pd.Series(filepaths, name='filepaths')
    Lseries=pd.Series(labels, name='labels')
    df=pd.concat([Fseries, Lseries], axis=1)            
    dsplit=vsplit/(1-trsplit)
    strat=df['labels']
    train_df, dummy_df=train_test_split(df, train_size=trsplit, shuffle=True, random_state=123, stratify=strat)
    strat=dummy_df['labels']
    valid_df, test_df= train_test_split(dummy_df, train_size=dsplit, shuffle=True, random_state=123, stratify=strat)
    print('train_df length: ', len(train_df), '  test_df length: ',len(test_df), '  valid_df length: ', len(valid_df))
    print(list(train_df['labels'].value_counts()))
    return train_df, test_df, valid_df

现在调用函数

sdir=r'C:\sdir'
trsplit=.8 # percent of images to use for training
vsplit=.1 # percent of images to use for validation
train_df, test_df, valid_df= preprocess(sdir,trsplit, vsplit)

现在您需要使用ImageDataGenerator创建3个生成器。flow_from_dataframe。文档为here。

channels=3
batch_size=20 # set batch size based on model complexity and sie of images
img_shape=(img_size[0], img_size[1], channels)
# calculate test_batch_size and test_steps so that test_batch_size X test_steps = number of test images
# this ensures you go through the test set exactly once when doing predictions on the test set
length=len(test_df)
test_batch_size=sorted([int(length/n) for n in range(1,length+1) if length % n ==0 and length/n<=80],reverse=True)[0]  
test_steps=int(length/test_batch_size)
print ( 'test batch size: ' ,test_batch_size, '  test steps: ', test_steps)
trgen=ImageDataGenerator(horizontal_flip=True)
tvgen=ImageDataGenerator()
msg='                                                              for the train generator'
print(msg, '\r', end='') 
train_gen=trgen.flow_from_dataframe( train_df, x_col='filepaths', y_col='labels', target_size=img_size, class_mode='categorical',
                                    color_mode='rgb', shuffle=True, batch_size=batch_size)
msg='                                                              for the test generator'
print(msg, '\r', end='') 
test_gen=tvgen.flow_from_dataframe( test_df, x_col='filepaths', y_col='labels', target_size=img_size, class_mode='categorical',
                                    color_mode='rgb', shuffle=False, batch_size=test_batch_size)
msg='                                                             for the validation generator'
print(msg, '\r', end='')
valid_gen=tvgen.flow_from_dataframe( valid_df, x_col='filepaths', y_col='labels', target_size=img_size, class_mode='categorical',
                                    color_mode='rgb', shuffle=True, batch_size=batch_size)
classes=list(train_gen.class_indices.keys())
class_count=len(classes)
train_steps=int(np.ceil(len(train_gen.labels)/batch_size))
labels=test_gen.labels

现在创建您的模型。下面显示了一个使用EfficientNetB3的建议模型

def make_model(img_img_size, class_count,lr=.001, trainable=True):
    img_shape=(img_size[0], img_size[1], 3)
    model_name='EfficientNetB3'
    base_model=tf.keras.applications.efficientnet.EfficientNetB3(include_top=False, weights="imagenet",input_shape=img_shape, pooling='max') 
    base_model.trainable=trainable
    x=base_model.output
    x=keras.layers.BatchNormalization(axis=-1, momentum=0.99, epsilon=0.001 )(x)
    x = Dense(256, kernel_regularizer = regularizers.l2(l = 0.016),activity_regularizer=regularizers.l1(0.006),
                    bias_regularizer=regularizers.l1(0.006) ,activation='relu')(x)
    x=Dropout(rate=.45, seed=123)(x)        
    output=Dense(class_count, activation='softmax')(x)
    model=Model(inputs=base_model.input, outputs=output)
    model.compile(Adamax(learning_rate=lr), loss='categorical_crossentropy', metrics=['accuracy']) 
    return model, base_model # return the base_model so the callback can control its training state

现在调用函数

model, base_model=make_model(img_size, class_count)

现在您可以训练模型

history=model.fit(x=train_gen,  epochs=epochs, verbose=0, validation_data=valid_gen,
               validation_steps=None,  shuffle=False,  initial_epoch=0)

训练后，您可以在测试集上评估模型性能

loss, acc=model.evaluate(test_gen, steps=test_steps)