1

2

mkdir pretrain_checkpoint

wget -P ./pretrain_checkpoint

1

2

3

git clone https://gitee.com/mindspore/mindspore.git -b r1.0

mkdir mobilenetv2

cp -r ./mindspore/model_zoo/official/cv/mobilenetv2 ./mobilenetv2/code

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

class MobileNetV2Backbone(nn.Cell):

   """

  MobileNetV2 architecture.

  Args:

      class_num (int): number of classes.

      width_mult (int): Channels multiplier for round to 8/16 and others. Default is 1.

      has_dropout (bool): Is dropout used. Default is false

      inverted_residual_setting (list): Inverted residual settings. Default is None

      round_nearest (list): Channel round to . Default is 8

  Returns:

      Tensor, output tensor.

  Examples:

      >>> MobileNetV2(num_classes=1000)

  """

   def __init__(self, width_mult=1., inverted_residual_setting=None, round_nearest=8,

                input_channel=32, last_channel=1280):

       super(MobileNetV2Backbone, self).__init__()

       block = InvertedResidual

       # setting of inverted residual blocks

       self.cfgs = inverted_residual_setting

       if inverted_residual_setting is None:

           self.cfgs = [

               # t, c, n, s

              [1, 16, 1, 1],

              [6, 24, 2, 2],

              [6, 32, 3, 2],

              [6, 64, 4, 2],

              [6, 96, 3, 1],

              [6, 160, 3, 2],

              [6, 320, 1, 1],

          ]

       # building first layer

       input_channel = _make_divisible(input_channel * width_mult, round_nearest)

       self.out_channels = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)

       features = [ConvBNReLU(3, input_channel, stride=2)]

       # building inverted residual blocks

       for t, c, n, s in self.cfgs:

           output_channel = _make_divisible(c * width_mult, round_nearest)

           for i in range(n):

               stride = s if i == 0 else 1

               features.append(block(input_channel, output_channel, stride, expand_ratio=t))

               input_channel = output_channel

       # building last several layers

       features.append(ConvBNReLU(input_channel, self.out_channels, kernel_size=1))

       # make it nn.CellList

       self.features = nn.SequentialCell(features)

       self._initialize_weights()

   def construct(self, x):

       x = self.features(x)

       return x

   def _initialize_weights(self):

       """

      Initialize weights.

      Args:

      Returns:

          None.

      Examples:

          >>> _initialize_weights()

      """

       self.init_parameters_data()

       for _, m in self.cells_and_names():

           if isinstance(m, nn.Conv2d):

               n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels

               m.weight.set_data(Tensor(np.random.normal(0, np.sqrt(2. / n),

                                                         m.weight.data.shape).astype("float32")))

               if m.bias is not None:

                   m.bias.set_data(

                       Tensor(np.zeros(m.bias.data.shape, dtype="float32")))

           elif isinstance(m, nn.BatchNorm2d):

               m.gamma.set_data(

                   Tensor(np.ones(m.gamma.data.shape, dtype="float32")))

               m.beta.set_data(

                   Tensor(np.zeros(m.beta.data.shape, dtype="float32")))

   @property

   def get_features(self):

       return self.features

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

class MobileNetV2Head(nn.Cell):

   """

  MobileNetV2 architecture.

  Args:

      class_num (int): Number of classes. Default is 1000.

      has_dropout (bool): Is dropout used. Default is false

  Returns:

      Tensor, output tensor.

  Examples:

      >>> MobileNetV2(num_classes=1000)

  """

   def __init__(self, input_channel=1280, num_classes=1000, has_dropout=False, activation="None"):

       super(MobileNetV2Head, self).__init__()

       # mobilenet head

       head = ([GlobalAvgPooling(), nn.Dense(input_channel, num_classes, has_bias=True)]

               if not has_dropout else

              [GlobalAvgPooling(), nn.Dropout(0.2), nn.Dense(input_channel, num_classes, has_bias=True)])

       self.head = nn.SequentialCell(head)

       self.need_activation = True

       if activation == "Sigmoid":

           self.activation = P.Sigmoid()

       elif activation == "Softmax":

           self.activation = P.Softmax()

       else:

           self.need_activation = False

       self._initialize_weights()

   def construct(self, x):

       x = self.head(x)

       if self.need_activation:

           x = self.activation(x)

       return x

   def _initialize_weights(self):

       """

      Initialize weights.

      Args:

      Returns:

          None.

      Examples:

          >>> _initialize_weights()

      """

       self.init_parameters_data()

       for _, m in self.cells_and_names():

           if isinstance(m, nn.Dense):

               m.weight.set_data(Tensor(np.random.normal(

                   0, 0.01, m.weight.data.shape).astype("float32")))

               if m.bias is not None:

                   m.bias.set_data(

                       Tensor(np.zeros(m.bias.data.shape, dtype="float32")))

   @property

   def get_head(self):

       return self.head

1

2

3

4

5

6

7

8

def define_net(config, is_training):

   backbone_net = MobileNetV2Backbone()

   activation = config.activation if not is_training else "None"

   head_net = MobileNetV2Head(input_channel=backbone_net.out_channels,

                              num_classes=config.num_classes,

                              activation=activation)

   net = mobilenet_v2(backbone_net, head_net)

   return backbone_net, head_net, net

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

def extract_features(net, dataset_path, config):

   features_folder = os.path.abspath(dataset_path) + '_features'

   if not os.path.exists(features_folder):

       os.makedirs(features_folder)

   dataset = create_dataset(dataset_path=dataset_path,

                            do_train=False,

                            config=config)

   step_size = dataset.get_dataset_size()

   if step_size == 0:

       raise ValueError("The step_size of dataset is zero. Check if the images count of train dataset \

                        is more than batch_size in config.py")

   model = Model(net)

   for i, data in enumerate(dataset.create_dict_iterator(output_numpy=True)):

       features_path = os.path.join(features_folder, f"feature_{i}.npy")

       label_path = os.path.join(features_folder, f"label_{i}.npy")

       if not os.path.exists(features_path) or not os.path.exists(label_path):

           image = data["image"]

           label = data["label"]

           features = model.predict(Tensor(image))

           np.save(features_path, features.asnumpy())

           np.save(label_path, label)

       print(f"Complete the batch {i+1}/{step_size}")

   return step_size

1

2

3

4

5

6

# define network

backbone_net, head_net, net = define_net(config, args_opt.is_training)

if args_opt.pretrain_ckpt and args_opt.freeze_layer == "backbone":

   load_ckpt(backbone_net, args_opt.pretrain_ckpt, trainable=False)

   step_size = extract_features(backbone_net, args_opt.dataset_path, config)

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

if args_opt.pretrain_ckpt is None or args_opt.freeze_layer == "none":

   loss_scale = FixedLossScaleManager(config.loss_scale, drop_overflow_update=False)

   opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr, config.momentum, \

           config.weight_decay, config.loss_scale)

   model = Model(net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale)

   cb = config_ckpoint(config, lr, step_size)

   print("============== Starting Training ==============")

   model.train(epoch_size, dataset, callbacks=cb)

   print("============== End Training ==============")

else:

   opt = Momentum(filter(lambda x: x.requires_grad, head_net.get_parameters()), lr, config.momentum,\      

                  config.weight_decay)

   network = WithLossCell(head_net, loss)

   network = TrainOneStepCell(network, opt)

   network.set_train()

   features_path = args_opt.dataset_path + '_features'

   idx_list = list(range(step_size))

   rank = 0

   if config.run_distribute:

       rank = get_rank()

   save_ckpt_path = os.path.join(config.save_checkpoint_path, 'ckpt_' + str(rank) + '/')

   if not os.path.isdir(save_ckpt_path):

       os.mkdir(save_ckpt_path)

   for epoch in range(epoch_size):

       random.shuffle(idx_list)

       epoch_start = time.time()

       losses = []

       for j in idx_list:

           feature = Tensor(np.load(os.path.join(features_path, f"feature_{j}.npy")))

           label = Tensor(np.load(os.path.join(features_path, f"label_{j}.npy")))

           losses.append(network(feature, label).asnumpy())

       epoch_mseconds = (time.time()-epoch_start) * 1000

       per_step_mseconds = epoch_mseconds / step_size

       print("epoch[{}/{}], iter[{}] cost: {:5.3f}, per step time: {:5.3f}, avg loss: {:5.3f}"\

      .format(epoch + 1, epoch_size, step_size, epoch_mseconds, per_step_mseconds,

               np.mean(np.array(losses))))

       if (epoch + 1) % config.save_checkpoint_epochs == 0:

           save_checkpoint(net, os.path.join(save_ckpt_path, f"mobilenetv2_{epoch+1}.ckpt"))

   print("total cost {:5.4f} s".format(time.time() - start))

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

config_cpu = ed({

   "num_classes": 5,        # 类别数,根据实验数据集调整

   "image_height": 224,     # 参与训练的图片高度

   "image_width": 224,      # 参与训练的图片宽度

   "batch_size": 64,        # 每批次样本数

   "epoch_size": 20,        # 训练次数

   "warmup_epochs": 0,      # 学习率变换步数

   "lr_init": .0,           # 初始化学习率

   "lr_end": 0.0,           # 最终学习率

   "lr_max": 0.01,          # 最大学习率

   "momentum": 0.9,         # momentum优化器参数

   "weight_decay": 4e-5,    # momentum优化器正则化参数

   "label_smooth": 0.1,     # 损失函数参数:smooth factor

   "loss_scale": 1024,      # "none"方式下损失函数参数,"backbone"方式下无用

   "save_checkpoint": True, # 是否保存模型

   "save_checkpoint_epochs": 5,    # "backbone"方式下,多少个epochs保存一次模型,"none"方式下无用

   "keep_checkpoint_max": 20,      # "none"方式下最多保存多少个模型,"backbone"方式下无用.

   "save_checkpoint_path": "./",   # 保存模型路径

   "platform": args.platform,      # 平台.可选"CPU", "GPU", "Ascend",这里固定"CPU"

   "run_distribute": False,        # 是否并行执行，"CPU"环境下默认False

   "activation": "Softmax"         # 最后一层激活函数

})

1

2

3

4

5

6

7

8

9

10

11

12

13

def train_parse_args():

   train_parser = argparse.ArgumentParser(description='Image classification trian')

   train_parser.add_argument('--platform', type=str, default="CPU", choices=("CPU", "GPU", "Ascend"), \

       help='run platform, only support CPU, GPU and Ascend')

   train_parser.add_argument('--dataset_path', type=str, default="../data/flower_photos_train", help='Dataset path')

   train_parser.add_argument('--pretrain_ckpt', type=str, default="../pretrain_checkpoint/mobilenetv2_cpu_gpu.ckpt", help='Pretrained checkpoint path \

      for fine tune or incremental learning')

   train_parser.add_argument('--freeze_layer', type=str, default="backbone", choices=["none", "backbone"], \

       help="freeze the weights of network from start to which layers")

   train_parser.add_argument('--run_distribute', type=ast.literal_, default=True, help='Run distribute')

   train_args = train_parser.parse_args()

   train_args.is_training = True

   return train_args

1

python train.py --platform CPU --dataset_path ../data/flower_photos_train --freeze_layer backbone --pretrain_ckpt ../pretrain_checkpoint/mobilenetv2_cpu_gpu.ckpt --run_distribute False

1

python .py --platform CPU --dataset_path ../data/flower_photos_test  --pretrain_ckpt ./ckpt_0/mobilenetv2_{epoch_num}.ckpt  --run_distribute False