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_criteo.""" 16import os 17import pytest 18 19from mindspore import context 20from mindspore.train.model import Model 21from mindspore.common import set_seed 22 23from src.deepfm import ModelBuilder, AUCMetric 24from src.config import DataConfig, ModelConfig, TrainConfig 25from src.dataset import create_dataset, DataType 26from src.callback import EvalCallBack, LossCallBack, TimeMonitor 27 28set_seed(1) 29 30@pytest.mark.level0 31@pytest.mark.platform_arm_ascend_training 32@pytest.mark.platform_x86_ascend_training 33@pytest.mark.env_onecard 34def test_deepfm(): 35 data_config = DataConfig() 36 train_config = TrainConfig() 37 device_id = int(os.getenv('DEVICE_ID')) 38 context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=device_id) 39 rank_size = None 40 rank_id = None 41 42 dataset_path = "/home/workspace/mindspore_dataset/criteo_data/mindrecord/" 43 print("dataset_path:", dataset_path) 44 ds_train = create_dataset(dataset_path, 45 train_mode=True, 46 epochs=1, 47 batch_size=train_config.batch_size, 48 data_type=DataType(data_config.data_format), 49 rank_size=rank_size, 50 rank_id=rank_id) 51 52 model_builder = ModelBuilder(ModelConfig, TrainConfig) 53 train_net, eval_net = model_builder.get_train_eval_net() 54 auc_metric = AUCMetric() 55 model = Model(train_net, eval_network=eval_net, metrics={"auc": auc_metric}) 56 57 loss_file_name = './loss.log' 58 time_callback = TimeMonitor(data_size=ds_train.get_dataset_size()) 59 loss_callback = LossCallBack(loss_file_path=loss_file_name) 60 callback_list = [time_callback, loss_callback] 61 62 eval_file_name = './auc.log' 63 ds_eval = create_dataset(dataset_path, train_mode=False, 64 epochs=1, 65 batch_size=train_config.batch_size, 66 data_type=DataType(data_config.data_format)) 67 eval_callback = EvalCallBack(model, ds_eval, auc_metric, 68 eval_file_path=eval_file_name) 69 callback_list.append(eval_callback) 70 71 print("train_config.train_epochs:", train_config.train_epochs) 72 model.train(train_config.train_epochs, ds_train, callbacks=callback_list) 73 74 export_loss_value = 0.52 75 print("loss_callback.loss:", loss_callback.loss) 76 assert loss_callback.loss < export_loss_value 77 export_per_step_time = 30.0 78 print("time_callback:", time_callback.per_step_time) 79 assert time_callback.per_step_time < export_per_step_time 80 print("*******test case pass!********") 81