前言:
偶然看到一个垃圾分类的文章,感觉很有趣,利用作者开源的数据集训练一个用于垃圾分类的ResNet50网络,回顾一下网络结构熟悉迁移学习的思想。
ResNet50的结构:
ResNet50是由大量的bottleNeck block组成的,然后根据shortcut分支也叫skip-connect分支上有无卷积操作分为identityBlock,convBlock,当然你也可以叫别的名字只要清楚他们的区别就行。
ResNet50的四个block分布情况:
(1)3 一个convBlock + 两个identityBlock (convBlock的strides=1)
(2)4 一个convBlock + 三个identityBlock (convBlock的strides=2)
(3)6 一个convBlock + 五个identityBlock (convBlock的strides=2)
(4)3 一个convBlock + 两个个identityBlock (convBlock的strides=2)
小细节:当convBlock的strides=1时,经过该block的featuremaps尺寸不变,strides=2时变为w,h各为输入的一半
关于convBlock以及identityBlock的具体实现可以参考我的另一篇文章:
这篇文章中给出了基于keras的实现。
数据
数据是任何一个深度学习项目的灵魂,高质量的数据集是取得可靠模型的基础。
下载链接:https://github.com/garythung/trashnet/blob/master/data/dataset-resized.zip
下载好之后解压缩然后放在特定的目录下备用。
解压之后得到一个文件夹:dataset-resized/
这个文件夹之下又有6个子文件夹对应6种类别:cardboard、glass、metal、paper、plastic、trash。
由于总的数据量较少,不到3000张所以从头训练模型显然不是一个好方法,这里要做迁移学习。
数据样例:
训练
这里只有6种类别,所以要对原始的ResNet50的分类层做一个调整。保留所有的特征提取层,替换后续的分类层然后利用解压好的数据训练模型,大概在7-8epoch后模型基本收敛,另外为了防止过拟合可以在训练阶段做数据增强以及dropout操作。下面是搭建训练模型的代码,基于keras实现。
def ResNet50(input_tensor=None):
"""
Input: input image shape e.g.(224,224,3)
return: tensor(1,1,2048)
"""
if K.image_data_format() == 'channels_last':
bn_axis = 3
else:
bn_axis = 1
x = ZeroPadding2D((3,3))(input_tensor)
x = Conv2D(64,(7,7),strides=(2,2),name='conv1')(x)
x = BatchNormalization(axis=bn_axis,name='bn_conv1')(x)
x = Activation('relu')(x)
x = MaxPooling2D((3,3),strides=(2,2))(x)
# w = w/4 h = h/4
x = conv_block(x,3,[64,64,256],stage=2,block='a',strides=(1,1))
x = identi_block(x,3,[64,64,256],stage=2,block='b')
x = identi_block(x,3,[64,64,256],stage=2,block='c')
## no change
x = conv_block(x,3,[128,128,512],stage=3,block='a')
x = identi_block(x,3,[128,128,512],stage=3,block='b')
x = identi_block(x,3,[128,128,512],stage=3,block='c')
x = identi_block(x,3,[128,128,512],stage=3,block='d')
# w = w/8 h = h/8
x = conv_block(x,3,[256,256,1024],stage=4,block='a')
x = identi_block(x,3,[256,256,1024],stage=4,block='b')
x = identi_block(x,3,[256,256,1024],stage=4,block='c')
x = identi_block(x,3,[256,256,1024],stage=4,block='d')
x = identi_block(x,3,[256,256,1024],stage=4,block='e')
x = identi_block(x,3,[256,256,1024],stage=4,block='f')
# w = w/16. h=h/16
x = conv_block(x,3,[512,512,2048],stage=5,block='a')
x = identi_block(x,3,[512,512,2048],stage=5,block='b')
x = identi_block(x,3,[512,512,2048],stage=5,block='c')
## w/=32. h/=32. (7,7,2048). ---->>>>(224,224,3)
x = AveragePooling2D((7,7),name='avg_pool')(x)
return x
然后做类别调整和fine-tuning
def bulid_model(input_shape,dropout,fc_layers,num_classes):
inputs = Input(shape=input_shape,name='input_1')
x = ResNet50(input_tensor=inputs)
x = Flatten()(x)
for fc in fc_layers:
x = Dense(fc,activation='relu')(x)
#x = Dropout(dropout)(x)
predictions = Dense(num_classes,activation='softmax')(x)
model = Model(inputs=inputs,outputs=predictions)
model.summary()
return model
利用keras的api读取训练数据并做数据增强
def gen(train_path):
train_datagen = ImageDataGenerator(rotation_range=90,horizontal_flip=True,vertical_flip=True,fill_mode='nearest')
train_generator = train_datagen.flow_from_directory(directory=train_path,target_size=(img_w,img_h),batch_size=batch_size,class_mode='categorical')
return train_generator
由于要做迁移学习,所以还需要在ImageNet上训练好的ResNet50模型参数
try:
pre_trained_weights = 'resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5'
model.load_weights(pre_trained_weights,by_name=True)
except Exception as e:
print('load pre-trained weights error {}'.format(e))
指定checkpoint存储路径,选择Adam优化函数定义提前终止训练的条件然后开始训练
checkpoint = ModelCheckpoint(filepath='weights/weights-{epoch:03d}-{loss:.2f}.h5',monitor='loss',save_best_only=False,save_weights_only=True)
checkpoint.set_model(model)
model.compile(optimizer=Adam(lr=1e-5),loss='categorical_crossentropy',metrics=['accuracy'])
lr_reducer = ReduceLROnPlateau(monitor='loss',factor=np.sqrt(0.1),cooldown=0,patience=2,min_lr=0.5e-6)
earlystopping = EarlyStopping(monitor='loss',patience=5,verbose=1)
tensorbord = TensorBoard(log_dir='weights/logs',write_graph=True)
model.fit_generator(generator=train_generator,steps_per_epoch=1000,epochs=epochs,initial_epoch=0,callbacks=[checkpoint,lr_reducer,earlystopping,tensorbord])
由于在训练时图像被统一缩放到了224*224大小,所以训练速度很快,2个小时左右可以跑玩30个epoch,但是一般8个epoch时模型就已经收敛了基本没有提升,所以可以提前终止。