diff --git a/config.py b/config.py index 4243c5d..4ab75e8 100644 --- a/config.py +++ b/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. ' diff --git a/datasets/create_dataset.py b/datasets/create_dataset.py index 371596c..2db4492 100644 --- a/datasets/create_dataset.py +++ b/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( diff --git a/demo.py b/demo.py index 126cbd3..5d61918 100644 --- a/demo.py +++ b/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: diff --git a/expand_images/predict_download_image.py b/expand_images/predict_download_image.py index d9ddb03..a996b2b 100644 --- a/expand_images/predict_download_image.py +++ b/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() diff --git a/models/build_model.py b/models/build_model.py index 0874d14..1fd8dfd 100644 --- a/models/build_model.py +++ b/models/build_model.py @@ -5,6 +5,25 @@ 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: """准备模型和优化器 """ @@ -12,16 +31,19 @@ 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, diff --git a/online-service/model/customize_service.py b/online-service/model/customize_service.py index 88bc377..0a0118d 100644 --- a/online-service/model/customize_service.py +++ b/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') diff --git a/online-service/model/deploy_models/build_model.py b/online-service/model/deploy_models/build_model.py index 4b4a9bc..7c6cf5e 100644 --- a/online-service/model/deploy_models/build_model.py +++ b/online-service/model/deploy_models/build_model.py @@ -4,6 +4,25 @@ 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: """准备模型和优化器 """ @@ -11,16 +30,19 @@ 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, diff --git a/online-service/offline_service.py b/online-service/offline_service.py index ca1e227..495f3bc 100644 --- a/online-service/offline_service.py +++ b/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') diff --git a/train_classifier.py b/train_classifier.py index b249ebf..8919cb6 100644 --- a/train_classifier.py +++ b/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) diff --git a/train_classifier_online.py b/train_classifier_online.py index 50a0c94..8b234e2 100644 --- a/train_classifier_online.py +++ b/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, diff --git a/utils/find_lr.py b/utils/find_lr.py index ef54111..5b1ddc0 100644 --- a/utils/find_lr.py +++ b/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)