add coco and cityscapes pretrain weights for seg

82f132f1 · FlyingQianMM · ee2d40d4 · 82f132f1 · 82f132f1 · 82f132f1
11 changed file
--- a/docs/apis/models/detection.md
+++ b/docs/apis/models/detection.md
@@ -42,7 +42,7 @@ train(self, num_epochs, train_dataset, train_batch_size=8, eval_dataset=None, sa
 > > - **save_interval_epochs** (int): 模型保存间隔（单位：迭代轮数）。默认为20。
 > > - **log_interval_steps** (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认值为'output'。
-> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为None。
+> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，则自动下载在COCO数据集上预训练的模型权重；若为None，则不使用预训练模型。默认为None。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认优化器：fluid.layers.piecewise_decay衰减策略，fluid.optimizer.Momentum优化方法。
 > > - **learning_rate** (float): 默认优化器的学习率。默认为1.0/8000。
 > > - **warmup_steps** (int):  默认优化器进行warmup过程的步数。默认为1000。
@@ -129,7 +129,7 @@ train(self, num_epochs, train_dataset, train_batch_size=2, eval_dataset=None, sa
 > > - **save_interval_epochs** (int): 模型保存间隔（单位：迭代轮数）。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认值为'output'。
-> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为None。
+> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，则自动下载在COCO数据集上预训练的模型权重；为None，则不使用预训练模型。默认为None。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认优化器：fluid.layers.piecewise_decay衰减策略，fluid.optimizer.Momentum优化方法。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认为0.0025。
 > > - **warmup_steps** (int):  默认优化器进行warmup过程的步数。默认为500。

--- a/docs/apis/models/instance_segmentation.md
+++ b/docs/apis/models/instance_segmentation.md
@@ -34,7 +34,7 @@ train(self, num_epochs, train_dataset, train_batch_size=1, eval_dataset=None, sa
 > > - **save_interval_epochs** (int): 模型保存间隔（单位：迭代轮数）。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认值为'output'。
-> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为None。
+> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，则自动下载在COCO数据集上预训练的模型权重；若为None，则不使用预训练模型。默认为None。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认优化器：fluid.layers.piecewise_decay衰减策略，fluid.optimizer.Momentum优化方法。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认为0.00125。
 > > - **warmup_steps** (int):  默认优化器进行warmup过程的步数。默认为500。

--- a/docs/apis/models/semantic_segmentation.md
+++ b/docs/apis/models/semantic_segmentation.md
@@ -40,7 +40,7 @@ train(self, num_epochs, train_dataset, train_batch_size=2, eval_dataset=None, ev
 > > - **save_interval_epochs** (int): 模型保存间隔（单位：迭代轮数）。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认'output'
-> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认'IMAGENET'。
+> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，则自动下载在COCO数据集上预训练的模型权重；若为字符串'CITYSCAPES'，则自动下载在CITYSCAPES数据集上预训练的模型权重；为None，则不使用预训练模型。默认'IMAGENET'。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认的优化器：使用fluid.optimizer.Momentum优化方法，polynomial的学习率衰减策略。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认0.01。
 > > - **lr_decay_power** (float): 默认优化器学习率衰减指数。默认0.9。
@@ -209,7 +209,7 @@ train(self, num_epochs, train_dataset, train_batch_size=2, eval_dataset=None, ev
 > > - **save_interval_epochs** (int): 模型保存间隔（单位：迭代轮数）。默认为1。
 > > - **log_interval_steps** (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
 > > - **save_dir** (str): 模型保存路径。默认'output'
-> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet数据集上预训练的模型权重；若为None，则不使用预训练模型。默认'IMAGENET'。
+> > - **pretrain_weights** (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，则自动下载在ImageNet数据集上预训练的模型权重；若为字符串'CITYSCAPES'，则自动下载在CITYSCAPES图片数据上预训练的模型权重；为None，则不使用预训练模型。默认'IMAGENET'。
 > > - **optimizer** (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认的优化器：使用fluid.optimizer.Momentum优化方法，polynomial的学习率衰减策略。
 > > - **learning_rate** (float): 默认优化器的初始学习率。默认0.01。
 > > - **lr_decay_power** (float): 默认优化器学习率衰减指数。默认0.9。

--- a/paddlex/cv/models/base.py
+++ b/paddlex/cv/models/base.py
@@ -194,12 +194,36 @@ class BaseAPI:
                if os.path.exists(pretrain_dir):
                    os.remove(pretrain_dir)
                os.makedirs(pretrain_dir)
+            if pretrain_weights is not None and \
+                not os.path.isdir(pretrain_weights) \
+                and not os.path.isfile(pretrain_weights):
+                if self.model_type == 'classifier':
+                    if pretrain_weights not in ['IMAGENET']:
+                        logging.warning(
+                            "Pretrain_weights for classifier should be defined as directory path or parameter file or 'IMAGENET' or None, but it is {}, so we force to set it as 'IMAGENET'".
+                            format(pretrain_weights))
+                        pretrain_weights = 'IMAGENET'
+                elif self.model_type == 'detector':
+                    if pretrain_weights not in ['IMAGENET', 'COCO']:
+                        logging.warning(
+                            "Pretrain_weights for detector should be defined as directory path or parameter file or 'IMAGENET' or 'COCO' or None, but it is {}, so we force to set it as 'IMAGENET'".
+                            format(pretrain_weights))
+                        pretrain_weights = 'IMAGENET'
+                elif self.model_type == 'segmenter':
+                    if pretrain_weights not in [
+                            'IMAGENET', 'COCO', 'CITYSCAPES'
+                    ]:
+                        logging.warning(
+                            "Pretrain_weights for segmenter should be defined as directory path or parameter file or 'IMAGENET' or 'COCO' or 'CITYSCAPES', but it is {}, so we force to set it as 'IMAGENET'".
+                            format(pretrain_weights))
+                        pretrain_weights = 'IMAGENET'
            if hasattr(self, 'backbone'):
                backbone = self.backbone
            else:
                backbone = self.__class__.__name__
                if backbone == "HRNet":
                    backbone = backbone + "_W{}".format(self.width)
+            class_name = self.__class__.__name__
            pretrain_weights = get_pretrain_weights(
                pretrain_weights, class_name, backbone, pretrain_dir)
        if startup_prog is None:

--- a/paddlex/cv/models/deeplabv3p.py
+++ b/paddlex/cv/models/deeplabv3p.py
@@ -242,7 +242,9 @@ class DeepLabv3p(BaseAPI):
            log_interval_steps (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
            save_dir (str): 模型保存路径。默认'output'。
            pretrain_weights (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，
-                则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认'IMAGENET。
+                则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，
+                则自动下载在COCO数据集上预训练的模型权重；若为字符串'CITYSCAPES'，
+                则自动下载在CITYSCAPES数据集上预训练的模型权重；若为None，则不使用预训练模型。默认'IMAGENET。
            optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认的优化器：使用
                fluid.optimizer.Momentum优化方法，polynomial的学习率衰减策略。
            learning_rate (float): 默认优化器的初始学习率。默认0.01。

--- a/paddlex/cv/models/faster_rcnn.py
+++ b/paddlex/cv/models/faster_rcnn.py
@@ -188,7 +188,8 @@ class FasterRCNN(BaseAPI):
            log_interval_steps (int): 训练日志输出间隔（单位：迭代次数）。默认为20。
            save_dir (str): 模型保存路径。默认值为'output'。
            pretrain_weights (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，
-                则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为'IMAGENET'。
+                则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，
+                则自动下载在COCO数据集上预训练的模型权重；若为None，则不使用预训练模型。默认为'IMAGENET'。
            optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认优化器：
                fluid.layers.piecewise_decay衰减策略，fluid.optimizer.Momentum优化方法。
            learning_rate (float): 默认优化器的初始学习率。默认为0.0025。

--- a/paddlex/cv/models/hrnet.py
+++ b/paddlex/cv/models/hrnet.py
@@ -150,7 +150,8 @@ class HRNet(DeepLabv3p):
            log_interval_steps (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
            save_dir (str): 模型保存路径。默认'output'。
            pretrain_weights (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，
-                则自动下载在IMAGENET图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为'IMAGENET'。
+                则自动下载在IMAGENET图片数据上预训练的模型权重；若为字符串'CITYSCAPES'
+                则自动下载在CITYSCAPES图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为'IMAGENET'。
            optimizer (paddle.fluid.optimizer): 优化器。当改参数为None时，使用默认的优化器：使用
                fluid.optimizer.Momentum优化方法，polynomial的学习率衰减策略。
            learning_rate (float): 默认优化器的初始学习率。默认0.01。

--- a/paddlex/cv/models/mask_rcnn.py
+++ b/paddlex/cv/models/mask_rcnn.py
@@ -147,7 +147,8 @@ class MaskRCNN(FasterRCNN):
            log_interval_steps (int): 训练日志输出间隔（单位：迭代次数）。默认为20。
            save_dir (str): 模型保存路径。默认值为'output'。
            pretrain_weights (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，
-                则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为None。
+                则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，
+                则自动下载在COCO数据集上预训练的模型权重；若为None，则不使用预训练模型。默认为None。
            optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认优化器：
                fluid.layers.piecewise_decay衰减策略，fluid.optimizer.Momentum优化方法。
            learning_rate (float): 默认优化器的学习率。默认为1.0/800。

--- a/paddlex/cv/models/utils/pretrain_weights.py
+++ b/paddlex/cv/models/utils/pretrain_weights.py
@@ -102,7 +102,20 @@ coco_pretrain = {
    'https://paddlemodels.bj.bcebos.com/object_detection/faster_rcnn_r101_fpn_2x.tar',
    'MaskRCNN_ResNet101_vd':
    'https://paddlemodels.bj.bcebos.com/object_detection/mask_rcnn_r101_vd_fpn_1x.tar',
-    'UNet': 'https://paddleseg.bj.bcebos.com/models/unet_coco_v3.tgz'
+    'UNet': 'https://paddleseg.bj.bcebos.com/models/unet_coco_v3.tgz',
+    'DeepLabv3p_MobileNetV2_x1.0':
+    'https://bj.bcebos.com/v1/paddleseg/deeplab_mobilenet_x1_0_coco.tgz',
+    'DeepLabv3p_Xception65':
+    'https://paddleseg.bj.bcebos.com/models/xception65_coco.tgz'
+}
+
+cityscapes_pretrain = {
+    'DeepLabv3p_MobileNetV2_x1.0':
+    'https://paddleseg.bj.bcebos.com/models/mobilenet_cityscapes.tgz',
+    'DeepLabv3p_Xception65':
+    'https://paddleseg.bj.bcebos.com/models/xception65_bn_cityscapes.tgz',
+    'HRNet_W18':
+    'https://paddleseg.bj.bcebos.com/models/hrnet_w18_bn_cityscapes.tgz'
 }


@@ -111,19 +124,38 @@ def get_pretrain_weights(flag, class_name, backbone, save_dir):
        return None
    elif osp.isdir(flag):
        return flag
-    warning_info = "{} supports to be finetuned with weights pretrained on the IMAGENET dataset only, so pretrain_weights is forced to be set to IMAGENET"
+    elif osp.isfile(flag):
+        return flag
+    warning_info = "{} does not support to be finetuned with weights pretrained on the {} dataset, so pretrain_weights is forced to be set to {}"
    if flag == 'COCO':
        if class_name == "FasterRCNN" and backbone in ['ResNet18'] or \
            class_name == "MaskRCNN" and backbone in ['ResNet18', 'HRNet_W18'] or \
            class_name == 'DeepLabv3p' and backbone in ['Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5', 'MobileNetV2_x1.5', 'MobileNetV2_x2.0']:
            model_name = '{}_{}'.format(class_name, backbone)
-            logging.warning(warning_info.format(model_name))
+            logging.warning(warning_info.format(model_name, flag, 'IMAGENET'))
            flag = 'IMAGENET'
        elif class_name == 'HRNet':
-            logging.warning(warning_info.format(class_name))
+            logging.warning(warning_info.format(class_name, flag, 'IMAGENET'))
            flag = 'IMAGENET'
-    if flag == 'CITYSCAPES':
+    elif flag == 'CITYSCAPES':
        model_name = '{}_{}'.format(class_name, backbone)
+        if class_name == 'UNet':
+            logging.warning(warning_info.format(class_name, flag, 'IMAGENET'))
+            flag = 'COCO'
+        if class_name == 'HRNet' and backbone.split('_')[
+                -1] in ['W30', 'W32', 'W40', 'W48', 'W60', 'W64']:
+            logging.warning(warning_info.format(backbone, flag, 'IMAGENET'))
+            flag = 'IMAGENET'
+        if class_name == 'DeepLabv3p' and backbone in [
+                'Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5',
+                'MobileNetV2_x1.5', 'MobileNetV2_x2.0'
+        ]:
+            model_name = '{}_{}'.format(class_name, backbone)
+            logging.warning(warning_info.format(model_name, flag, 'IMAGENET'))
+            flag = 'IMAGENET'
+    elif flag == 'IMAGENET' and class_name == 'UNet':
+        logging.warning(warning_info.format(class_name, flag, 'COCO'))
+        flag = 'COCO'

    if flag == 'IMAGENET':
        new_save_dir = save_dir
@@ -160,19 +192,19 @@ def get_pretrain_weights(flag, class_name, backbone, save_dir):
                raise Exception(
                    "Unexpected error, please make sure paddlehub >= 1.6.2")
        return osp.join(new_save_dir, backbone)
-    elif flag == 'COCO':
+    elif flag in ['COCO', 'CITYSCAPES']:
        new_save_dir = save_dir
        if hasattr(paddlex, 'pretrain_dir'):
            new_save_dir = paddlex.pretrain_dir
-        if class_name in ['YOLOv3', 'FasterRCNN', 'MaskRCNN']:
+        if class_name in ['YOLOv3', 'FasterRCNN', 'MaskRCNN', 'DeepLabv3p']:
            backbone = '{}_{}'.format(class_name, backbone)
-        url = coco_pretrain[backbone]
+        if flag == 'COCO':
+            url = coco_pretrain[backbone]
+        elif flag == 'CITYSCAPES':
+            url = cityscapes_pretrain[backbone]
        fname = osp.split(url)[-1].split('.')[0]
        #        paddlex.utils.download_and_decompress(url, path=new_save_dir)
        #        return osp.join(new_save_dir, fname)
-
-        assert backbone in coco_pretrain, "There is not COCO pretrain weights for {}, you may try ImageNet.".format(
-            backbone)
        try:
            hub.download(backbone, save_path=new_save_dir)
        except Exception as e:
@@ -189,5 +221,5 @@ def get_pretrain_weights(flag, class_name, backbone, save_dir):
        return osp.join(new_save_dir, backbone)
    else:
        raise Exception(
-            "pretrain_weights need to be defined as directory path or `IMAGENET` or 'COCO' (download pretrain weights automatically)."
+            "pretrain_weights need to be defined as directory path or 'IMAGENET' or 'COCO' or 'Cityscapes' (download pretrain weights automatically)."
        )
--- a/paddlex/cv/models/yolo_v3.py
+++ b/paddlex/cv/models/yolo_v3.py
@@ -180,7 +180,8 @@ class YOLOv3(BaseAPI):
            log_interval_steps (int): 训练日志输出间隔（单位：迭代次数）。默认为10。
            save_dir (str): 模型保存路径。默认值为'output'。
            pretrain_weights (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，
-                则自动下载在ImageNet图片数据上预训练的模型权重；若为None，则不使用预训练模型。默认为'IMAGENET'。
+                则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，
+                则自动下载在COCO数据集上预训练的模型权重；若为None，则不使用预训练模型。默认为'IMAGENET'。
            optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认优化器：
                fluid.layers.piecewise_decay衰减策略，fluid.optimizer.Momentum优化方法。
            learning_rate (float): 默认优化器的学习率。默认为1.0/8000。

--- a/paddlex/utils/utils.py
+++ b/paddlex/utils/utils.py
@@ -67,8 +67,8 @@ def parse_param_file(param_file, return_shape=True):
        f.close()
        return tuple(tensor_desc.dims)
    if tensor_desc.data_type != 5:
-        raise Exception(
-            "Unexpected data type while parse {}".format(param_file))
+        raise Exception("Unexpected data type while parse {}".format(
+            param_file))
    data_size = 4
    for i in range(len(tensor_shape)):
        data_size *= tensor_shape[i]
@@ -139,7 +139,12 @@ def load_pdparams(exe, main_prog, model_dir):

    vars_to_load = list()
    import pickle
-    with open(osp.join(model_dir, 'model.pdparams'), 'rb') as f:
+
+    if osp.isfile(model_dir):
+        params_file = model_dir
+    else:
+        params_file = osp.join(model_dir, 'model.pdparams')
+    with open(params_file, 'rb') as f:
        params_dict = pickle.load(f) if six.PY2 else pickle.load(
            f, encoding='latin1')
    unused_vars = list()
@@ -185,8 +190,8 @@ def is_belong_to_optimizer(var):
    import paddle.fluid as fluid
    from paddle.fluid.proto.framework_pb2 import VarType

-    if not (isinstance(var, fluid.framework.Parameter)
-            or var.desc.need_check_feed()):
+    if not (isinstance(var, fluid.framework.Parameter) or
+            var.desc.need_check_feed()):
        return is_persistable(var)
    return False

@@ -206,9 +211,8 @@ def load_pdopt(exe, main_prog, model_dir):
    if len(optimizer_var_list) > 0:
        for var in optimizer_var_list:
            if var.name not in opt_dict:
-                raise Exception(
-                    "{} is not in saved paddlex optimizer, {}".format(
-                        var.name, exception_message))
+                raise Exception("{} is not in saved paddlex optimizer, {}".
+                                format(var.name, exception_message))
            if var.shape != opt_dict[var.name].shape:
                raise Exception(
                    "Shape of optimizer variable {} doesn't match.(Last: {}, Now: {}), {}"
@@ -227,9 +231,8 @@ def load_pdopt(exe, main_prog, model_dir):
            "There is no optimizer parameters in the model, please set the optimizer!"
        )
    else:
-        logging.info(
-            "There are {} optimizer parameters in {} are loaded.".format(
-                len(optimizer_var_list), model_dir))
+        logging.info("There are {} optimizer parameters in {} are loaded.".
+                     format(len(optimizer_var_list), model_dir))


 def load_pretrain_weights(exe,
@@ -239,6 +242,12 @@ def load_pretrain_weights(exe,
                          resume=False):
    if not osp.exists(weights_dir):
        raise Exception("Path {} not exists.".format(weights_dir))
+    if osp.isfile(weights_dir):
+        if not weights_dir.endswith('.pdparams'):
+            raise Exception("File {} is not a paddle parameter file".format(
+                weights_dir))
+        load_pdparams(exe, main_prog, weights_dir)
+        return
    if osp.exists(osp.join(weights_dir, "model.pdparams")):
        load_pdparams(exe, main_prog, weights_dir)
        if resume:
@@ -255,9 +264,8 @@ def load_pretrain_weights(exe,
        if not isinstance(var, fluid.framework.Parameter):
            continue
        if not osp.exists(osp.join(weights_dir, var.name)):
-            logging.debug(
-                "[SKIP] Pretrained weight {}/{} doesn't exist".format(
-                    weights_dir, var.name))
+            logging.debug("[SKIP] Pretrained weight {}/{} doesn't exist".
+                          format(weights_dir, var.name))
            continue
        pretrained_shape = parse_param_file(osp.join(weights_dir, var.name))
        actual_shape = tuple(var.shape)
@@ -317,9 +325,8 @@ def load_pretrain_weights(exe,
                "There is no optimizer parameters in the model, please set the optimizer!"
            )
        else:
-            logging.info(
-                "There are {} optimizer parameters in {} are loaded.".format(
-                    len(optimizer_var_list), weights_dir))
+            logging.info("There are {} optimizer parameters in {} are loaded.".
+                         format(len(optimizer_var_list), weights_dir))


 class EarlyStop:
@@ -342,12 +349,12 @@ class EarlyStop:
            self.max = current_score
            return False
        else:
-            if (abs(self.score - current_score) < self.thresh
-                    or current_score < self.score):
+            if (abs(self.score - current_score) < self.thresh or
+                    current_score < self.score):
                self.counter += 1
                self.score = current_score
-                logging.debug(
-                    "EarlyStopping: %i / %i" % (self.counter, self.patience))
+                logging.debug("EarlyStopping: %i / %i" %
+                              (self.counter, self.patience))
                if self.counter >= self.patience:
                    logging.info("EarlyStopping: Stop training")
                    return True