train.py

# -*- coding: utf-8 -*-
"""
Created on Sun Apr 21 13:52:07 2019

@author: Administrator

train.py: 训练模型

"""

from Unet import Unet
import LoadBatches1D
import keras
from keras import optimizers
import warnings
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore")


def lr_schedule(epoch):
    # 训练网络时学习率衰减方案
    lr = 0.0001
    if epoch >= 50:
        lr = 0.00001
    print('Learning rate: ', lr)
    return lr


train_sigs_path = 'G:/ECG_UNet/train_sigs/'
train_segs_path = 'G:/ECG_UNet/train_labels/'
train_batch_size = 1
n_classes = 3
input_length = 1800
optimizer_name = optimizers.Adam(lr_schedule(0))
val_sigs_path = 'G:/ECG_UNet/val_sigs/'
val_segs_path = 'G:/ECG_UNet/val_labels/'
val_batch_size = 2

lr_scheduler = keras.callbacks.LearningRateScheduler(lr_schedule)

model = Unet(n_classes, input_length=input_length)

model.compile(loss='categorical_crossentropy',
              optimizer=optimizer_name,
              metrics=['accuracy'])

model.summary()

output_length = 1800

G = LoadBatches1D.SigSegmentationGenerator(train_sigs_path, train_segs_path, train_batch_size, n_classes, output_length)

G2 = LoadBatches1D.SigSegmentationGenerator(val_sigs_path, val_segs_path, val_batch_size, n_classes, output_length)

checkpointer = keras.callbacks.ModelCheckpoint(filepath='myNet.h5', monitor='val_acc', mode='max', save_best_only=True)

history = model.fit_generator(G, 500//train_batch_size, validation_data=G2, validation_steps=200, epochs=70,
                        callbacks=[checkpointer, lr_scheduler])

plt.figure()
plt.plot(history.history['acc'])
plt.plot(history.history['val_acc'])
plt.title('Model accuracy')
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
plt.legend(['Train', 'Test'], loc='upper left')
plt.grid(True)

plt.figure()
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Test'], loc='upper left')
plt.grid(True)