add gn

config.py

@@ -62,6 +62,8 @@ def get_classify_config():
     parser.add_argument('--model_type', type=str, default='se_resnext101_32x4d',
                         help='densenet201/efficientnet-b5/se_resnext101_32x4d')
     parser.add_argument('--drop_rate', type=float, default=0, help='dropout rate in classify module')
+    parser.add_argument('--bn_to_gn', type=str2bool, nargs='?', const=True, default=False, 
+                        help='change bn to gn')
     parser.add_argument('--restore', type=str, default='',
                         help='Load the weight file before training.'
                              'if it is equal to `last`, load the `model_best.pth` in the last modification folder. '

datasets/create_dataset.py

@@ -476,7 +476,7 @@ def get_dataloader_from_folder(
         image_size, 
         mean=mean, 
         std=std, 
-        multi_scale=multi_scale
+        multi_scale=False
         )
 
     train_dataloader = DataLoader(

demo.py

@@ -10,7 +10,7 @@
 from PIL import Image, ImageFont, ImageDraw
 from models.build_model import PrepareModel
 from config import get_classify_config
-from datasets.create_dataset import GetDataloader
+from datasets.create_dataset import GetDataloader, get_dataloader_from_folder
 
 
 class DemoResults(object):
@@ -28,6 +28,7 @@ def __init__(self, config, weight_path, label_json_path, fold, mean=(0.485, 0.45
         self.model_type = config.model_type
         self.classes_num = config.num_classes
         self.drop_rate = config.drop_rate
+        self.bn_to_gn = config.bn_to_gn
         self.image_size = config.image_size
         self.weight_path = weight_path
         self.fold = str(fold)
@@ -121,7 +122,7 @@ def __prepare__(self, label_json_path):
             label_dict: dict，类标名称与类标之间的对应关系
         """
         prepare_model = PrepareModel()
-        model = prepare_model.create_model(self.model_type, self.classes_num, self.drop_rate, pretrained=False)
+        model = prepare_model.create_model(self.model_type, self.classes_num, self.drop_rate, pretrained=False, bn_to_gn=self.bn_to_gn)
         model.load_state_dict(torch.load(self.weight_path)['state_dict'])
         print('Successfully Loaded from %s' % self.weight_path)
         model = model.cuda()
@@ -152,22 +153,36 @@ def __prepare__(self, label_json_path):
 
     # 先删除该目录下所有的文件，再建立该文件夹
     save_path = 'data/demo_data/results'
-    shutil.rmtree(save_path)
+    if os.path.exists(save_path):
+        shutil.rmtree(save_path)
     os.makedirs(save_path)
 
-    get_dataloader = GetDataloader(
-        data_root, 
-        folds_split=folds_split, 
-        test_size=test_size, 
-        choose_dataset=config.choose_dataset,
-        load_split_from_file=config.load_split_from_file
-        )
-    train_dataloaders, val_dataloaders, _, _ = get_dataloader.get_dataloader(
-        config.batch_size, 
-        config.image_size, 
-        mean, std,
-        transforms=transforms
-        )
+    if config.dataset_from_folder:
+        train_dataloaders, val_dataloaders, train_labels_number, _ = get_dataloader_from_folder(
+            data_root,
+            config.image_size,
+            transforms,
+            mean,
+            std,
+            config.batch_size,
+            config.multi_scale
+            )
+        train_dataloaders, val_dataloaders, train_labels_number_folds = [train_dataloaders], [val_dataloaders], [train_labels_number]
+    else:
+
+        get_dataloader = GetDataloader(
+            data_root,
+            folds_split=folds_split,
+            test_size=test_size,
+            choose_dataset=config.choose_dataset,
+            load_split_from_file=config.load_split_from_file
+            )
+        train_dataloaders, val_dataloaders, _, _ = get_dataloader.get_dataloader(
+            config.batch_size,
+            config.image_size,
+            mean, std,
+            transforms=transforms
+            )
 
     for fold_index, [train_loader, valid_loader] in enumerate(zip(train_dataloaders, val_dataloaders)):
         if fold_index in config.selected_fold:

expand_images/predict_download_image.py

@@ -130,7 +130,7 @@ def save_image_label(self, save_path, image_path, image_name, label, index, scor
 
     def __prepare__(self, label_json_path):
         prepare_model = PrepareModel()
-        model = prepare_model.create_model(self.model_type, self.classes_num, 0, pretrained=False)
+        model = prepare_model.create_model(self.model_type, self.classes_num, 0, pretrained=False, bn_to_gn=False)
         model.load_state_dict(torch.load(self.weight_path)['state_dict'])
         model = model.cuda()
         model.eval()

models/build_model.py

@@ -5,23 +5,45 @@
 from models.custom_attention_model import CustomLocalAttentionModel
 
 
+def convert_layers(model, layer_type_old, layer_type_new, convert_weights=False, num_groups=None):
+    for name, module in reversed(model._modules.items()):
+        if len(list(module.children())) > 0:
+            # recurse
+            model._modules[name] = convert_layers(module, layer_type_old, layer_type_new, convert_weights)
+
+        if type(module) == layer_type_old:
+            layer_old = module
+            layer_new = layer_type_new(module.num_features if num_groups is None else num_groups, module.num_features, module.eps, module.affine)
+
+            if convert_weights:
+                layer_new.weight = layer_old.weight
+                layer_new.bias = layer_old.bias
+
+            model._modules[name] = layer_new
+
+    return model
+
+
 class PrepareModel:
     """准备模型和优化器
     """
 
     def __init__(self):
         pass
 
-    def create_model(self, model_type, classes_num, drop_rate=0, pretrained=True):
+    def create_model(self, model_type, classes_num, drop_rate=0, pretrained=True, bn_to_gn=False):
         """创建模型
         Args:
             model_type: str, 模型类型
             classes_num: int, 类别数目
             drop_rate: float, 分类层中的drop out系数
             pretrained: bool, 是否使用预训练模型
+            bn_to_gn: bool, 是否使用gn替换bn
         """
         print('Creating model: {}'.format(model_type))
         model = CustomModel(model_type, classes_num, drop_rate=drop_rate, pretrained=pretrained)
+        if bn_to_gn:
+            convert_layers(model, torch.nn.BatchNorm2d, torch.nn.GroupNorm, True, num_groups=16)
         return model
 
     def create_local_attention_model(self, model_type, classes_num, last_stride=2, drop_rate=0,

online-service/model/customize_service.py

@@ -167,7 +167,7 @@ def __prepare(self):
         """准备模型
         """
         prepare_model = PrepareModel()
-        model = prepare_model.create_model('se_resnext101_32x4d', self.classes_num, drop_rate=0, pretrained=False)
+        model = prepare_model.create_model('se_resnext101_32x4d', self.classes_num, drop_rate=0, pretrained=False, bn_to_gn=False)
 
         if torch.cuda.is_available():
             logger.info('Using GPU for inference')

online-service/model/deploy_models/build_model.py

@@ -4,23 +4,45 @@
 from model.deploy_models.custom_attention_model import CustomLocalAttentionModel
 
 
+def convert_layers(model, layer_type_old, layer_type_new, convert_weights=False, num_groups=None):
+    for name, module in reversed(model._modules.items()):
+        if len(list(module.children())) > 0:
+            # recurse
+            model._modules[name] = convert_layers(module, layer_type_old, layer_type_new, convert_weights)
+
+        if type(module) == layer_type_old:
+            layer_old = module
+            layer_new = layer_type_new(module.num_features if num_groups is None else num_groups, module.num_features, module.eps, module.affine)
+
+            if convert_weights:
+                layer_new.weight = layer_old.weight
+                layer_new.bias = layer_old.bias
+
+            model._modules[name] = layer_new
+
+    return model
+
+
 class PrepareModel:
     """准备模型和优化器
     """
 
     def __init__(self):
         pass
 
-    def create_model(self, model_type, classes_num, drop_rate=0, pretrained=True):
+    def create_model(self, model_type, classes_num, drop_rate=0, pretrained=True, bn_to_gn=False):
         """创建模型
         Args:
             model_type: str, 模型类型
             classes_num: int, 类别数目
             drop_rate: float, 分类层中的drop out系数
             pretrained: bool, 是否使用预训练模型
+            bn_to_gn: bool, 是否使用gn替换bn
         """
         print('Creating model: {}'.format(model_type))
         model = CustomModel(model_type, classes_num, drop_rate=drop_rate, pretrained=pretrained)
+        if bn_to_gn:
+            convert_layers(model, torch.nn.BatchNorm2d, torch.nn.GroupNorm, True, num_groups=16)
         return model
 
     def create_local_attention_model(self, model_type, classes_num, last_stride=2, drop_rate=0,

online-service/offline_service.py

@@ -161,7 +161,7 @@ def __prepare(self):
         """准备模型
         """
         prepare_model = PrepareModel()
-        model = prepare_model.create_model('se_resnext101_32x4d', self.classes_num, drop_rate=0, pretrained=False)
+        model = prepare_model.create_model('se_resnext101_32x4d', self.classes_num, drop_rate=0, pretrained=False, bn_to_gn=False)
 
         print('Using CPU for inference')
         checkpoint = torch.load(self.model_path, map_location='cpu')

train_classifier.py

@@ -57,7 +57,8 @@ def __init__(self, config, fold, train_labels_number):
             model_type=config.model_type,
             classes_num=self.num_classes,
             drop_rate=config.drop_rate,
-            pretrained=True
+            pretrained=True,
+            bn_to_gn=config.bn_to_gn
         )
         if torch.cuda.is_available():
             self.model = torch.nn.DataParallel(self.model)

train_classifier_online.py

@@ -112,13 +112,25 @@ def __init__(self, config, fold, train_labels_number):
             print('use exist pre-trained model at: %s' %
                   (os.path.abspath('/home/work/.cache/torch/checkpoints/se_resnext101_32x4d-3b2fe3d8.pth')))
 
+        # 拷贝预训练权重
+        print("=> using pre-trained model '{}'".format(config.model_type))
+        if not mox.file.exists('/home/work/.cache/torch/checkpoints/efficientnet-b5-b6417697.pth'):
+            mox.file.copy(os.path.join(self.bucket_name, 'model_zoo/efficientnet-b5-b6417697.pth'),
+                          '/home/work/.cache/torch/checkpoints/efficientnet-b5-b6417697.pth')
+            print('copy pre-trained model from OBS to: %s success' %
+                  (os.path.abspath('/home/work/.cache/torch/checkpoints/efficientnet-b5-b6417697.pth')))
+        else:
+            print('use exist pre-trained model at: %s' %
+                  (os.path.abspath('/home/work/.cache/torch/checkpoints/efficientnet-b5-b6417697.pth')))
+
         # 加载模型
         prepare_model = PrepareModel()
         self.model = prepare_model.create_model(
             model_type=config.model_type,
             classes_num=self.num_classes,
             drop_rate=config.drop_rate,
-            pretrained=True
+            pretrained=True,
+            bn_to_gn=config.bn_to_gn
         )
         self.model = torch.nn.DataParallel(self.model).cuda()
 
@@ -349,7 +361,7 @@ def init_log(self):
 
     if config.dataset_from_folder:
         train_dataloaders, val_dataloaders, train_labels_number, _ = get_dataloader_from_folder(
-            data_root, 
+            config.data_local, 
             config.image_size, 
             transforms, 
             mean, 

utils/find_lr.py

@@ -31,7 +31,8 @@ def prepare(config, train_labels_number):
         model_type=config.model_type,
         classes_num=config.num_classes,
         drop_rate=config.drop_rate,
-        pretrained=True
+        pretrained=True,
+        bn_to_gn=config.bn_to_gn
     )
     if torch.cuda.is_available():
         model = torch.nn.DataParallel(model)