基于ResNet50网络的简单垃圾分类网络

前言：

偶然看到一个垃圾分类的文章，感觉很有趣，利用作者开源的数据集训练一个用于垃圾分类的ResNet50网络，回顾一下网络结构熟悉迁移学习的思想。

原始文章：How to build an image classifier for waste sorting link

有一点区别：原文用的是ResNet34，博主这里用的是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的具体实现可以参考我的另一篇文章：https://blog.csdn.net/zshluckydogs/article/details/100588677
这篇文章中给出了基于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时模型就已经收敛了基本没有提升，所以可以提前终止。

完整的项目地址：https://github.com/luckydog5/waste-classification.git
直接克隆这个地址：git clone https://github.com/luckydog5/waste-classification.git
参考模型可以从我的百度云网盘下载：链接:https://pan.baidu.com/s/1bP61cO6-la0g2OHCYVZR0A 密码:6m4n

如果GitHub上链接下不了数据，可通过百度网盘下载。

链接:https://pan.baidu.com/s/1euEVQ7FmRCaljqyJ8wA4PQ 密码:cohe

欢迎讨论和提bug。