• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""train."""
16import argparse
17import time
18import pytest
19import numpy as np
20from mindspore import context, Tensor
21from mindspore.nn.optim.momentum import Momentum
22from mindspore import Model
23from mindspore.train.callback import Callback
24from src.md_dataset import create_dataset
25from src.losses import OhemLoss
26from src.deeplabv3 import deeplabv3_resnet50
27from src.config import config
28
29context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
30#--train
31#--eval
32#  --Images
33#    --2008_001135.jpg
34#    --2008_001404.jpg
35#  --SegmentationClassRaw
36#    --2008_001135.png
37#    --2008_001404.png
38data_url = "/home/workspace/mindspore_dataset/voc/voc2012"
39class LossCallBack(Callback):
40    """
41    Monitor the loss in training.
42    Note:
43        if per_print_times is 0 do not print loss.
44    Args:
45        per_print_times (int): Print loss every times. Default: 1.
46    """
47    def __init__(self, data_size, per_print_times=1):
48        super(LossCallBack, self).__init__()
49        if not isinstance(per_print_times, int) or per_print_times < 0:
50            raise ValueError("print_step must be int and >= 0")
51        self.data_size = data_size
52        self._per_print_times = per_print_times
53        self.time = 1000
54        self.loss = 0
55    def epoch_begin(self, run_context):
56        self.epoch_time = time.time()
57    def step_end(self, run_context):
58        cb_params = run_context.original_args()
59        epoch_mseconds = (time.time() - self.epoch_time) * 1000
60        self.time = epoch_mseconds / self.data_size
61        self.loss = cb_params.net_outputs
62        print("epoch: {}, step: {}, outputs are {}".format(cb_params.cur_epoch_num, cb_params.cur_step_num,
63                                                           str(cb_params.net_outputs)))
64
65def model_fine_tune(train_net, fix_weight_layer):
66    train_net.init_parameters_data()
67    for para in train_net.trainable_params():
68        para.set_data(Tensor(np.ones(para.data.shape).astype(np.float32) * 0.02))
69        if fix_weight_layer in para.name:
70            para.requires_grad = False
71
72@pytest.mark.level0
73@pytest.mark.platform_arm_ascend_training
74@pytest.mark.platform_x86_ascend_training
75@pytest.mark.env_onecard
76def test_deeplabv3_1p():
77    start_time = time.time()
78    epoch_size = 100
79    args_opt = argparse.Namespace(base_size=513, crop_size=513, batch_size=2)
80    args_opt.base_size = config.crop_size
81    args_opt.crop_size = config.crop_size
82    args_opt.batch_size = config.batch_size
83    train_dataset = create_dataset(args_opt, data_url, 1, config.batch_size,
84                                   usage="eval")
85    dataset_size = train_dataset.get_dataset_size()
86    callback = LossCallBack(dataset_size)
87    net = deeplabv3_resnet50(config.seg_num_classes, [config.batch_size, 3, args_opt.crop_size, args_opt.crop_size],
88                             infer_scale_sizes=config.eval_scales, atrous_rates=config.atrous_rates,
89                             decoder_output_stride=config.decoder_output_stride, output_stride=config.output_stride,
90                             fine_tune_batch_norm=config.fine_tune_batch_norm, image_pyramid=config.image_pyramid)
91    net.set_train()
92    model_fine_tune(net, 'layer')
93    loss = OhemLoss(config.seg_num_classes, config.ignore_label)
94    opt = Momentum(filter(lambda x: 'beta' not in x.name and 'gamma' not in x.name and 'depth' not in x.name and 'bias' not in x.name, net.trainable_params()), learning_rate=config.learning_rate, momentum=config.momentum, weight_decay=config.weight_decay)
95    model = Model(net, loss, opt)
96    model.train(epoch_size, train_dataset, callback)
97    print(time.time() - start_time)
98    print("expect loss: ", callback.loss)
99    print("expect time: ", callback.time)
100    expect_loss = 0.92
101    expect_time = 50
102    assert callback.loss.asnumpy() <= expect_loss
103    assert callback.time <= expect_time
104