wehub-resource-sync
173 lines
8.0 KiB
Python
173 lines
8.0 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
|
|
from tensorboard.backend.event_processing import event_accumulator
|
|
from visualdl import LogReader
|
|
|
|
from paddlenlp.trainer import TrainerControl, TrainerState, TrainingArguments
|
|
from paddlenlp.trainer.integrations import (
|
|
SwanLabCallback,
|
|
TensorBoardCallback,
|
|
VisualDLCallback,
|
|
WandbCallback,
|
|
)
|
|
from tests.trainer.trainer_utils import RegressionModelConfig, RegressionPretrainedModel
|
|
|
|
|
|
class TestWandbCallback(unittest.TestCase):
|
|
def test_wandbcallback(self):
|
|
output_dir = tempfile.mkdtemp()
|
|
args = TrainingArguments(
|
|
output_dir=output_dir,
|
|
max_steps=200,
|
|
logging_steps=20,
|
|
run_name="test_wandbcallback",
|
|
logging_dir=output_dir,
|
|
)
|
|
state = TrainerState(trial_name="PaddleNLP")
|
|
control = TrainerControl()
|
|
config = RegressionModelConfig(a=1, b=1)
|
|
model = RegressionPretrainedModel(config)
|
|
os.environ["WANDB_LOG_MODEL"] = "checkpoint"
|
|
os.environ["WANDB_MODE"] = "offline"
|
|
wandbcallback = WandbCallback()
|
|
self.assertFalse(wandbcallback._initialized)
|
|
wandbcallback.on_train_begin(args, state, control)
|
|
self.assertTrue(wandbcallback._initialized)
|
|
self.assertEqual(wandbcallback._wandb.run.name, state.trial_name)
|
|
self.assertEqual(wandbcallback._wandb.run.group, args.run_name)
|
|
for global_step in range(args.max_steps):
|
|
state.global_step = global_step
|
|
if global_step % args.logging_steps == 0:
|
|
log = {"loss": 100 - 0.4 * global_step, "learning_rate": 0.1, "global_step": global_step}
|
|
wandbcallback.on_log(args, state, control, logs=log)
|
|
self.assertEqual(wandbcallback._wandb.run.summary["train/loss"], log["loss"])
|
|
self.assertEqual(wandbcallback._wandb.run.summary["train/learning_rate"], log["learning_rate"])
|
|
self.assertEqual(wandbcallback._wandb.run.summary["train/global_step"], log["global_step"])
|
|
wandbcallback.on_train_end(args, state, control, model=model)
|
|
wandbcallback._wandb.finish()
|
|
os.environ.pop("WANDB_LOG_MODEL", None)
|
|
os.environ.pop("WANDB_MODE", None)
|
|
shutil.rmtree(output_dir)
|
|
|
|
|
|
class TestSwanlabCallback(unittest.TestCase):
|
|
def test_swanlabcallback(self):
|
|
output_dir = tempfile.mkdtemp()
|
|
args = TrainingArguments(
|
|
output_dir=output_dir,
|
|
max_steps=200,
|
|
logging_steps=20,
|
|
run_name="test_swanlabcallback",
|
|
logging_dir=output_dir,
|
|
)
|
|
state = TrainerState(trial_name="PaddleNLP")
|
|
control = TrainerControl()
|
|
config = RegressionModelConfig(a=1, b=1)
|
|
model = RegressionPretrainedModel(config)
|
|
os.environ["SWANLAB_MODE"] = "disabled"
|
|
swanlabcallback = SwanLabCallback()
|
|
self.assertFalse(swanlabcallback._initialized)
|
|
swanlabcallback.on_train_begin(args, state, control)
|
|
self.assertTrue(swanlabcallback._initialized)
|
|
for global_step in range(args.max_steps):
|
|
state.global_step = global_step
|
|
if global_step % args.logging_steps == 0:
|
|
log = {"loss": 100 - 0.4 * global_step, "learning_rate": 0.1, "global_step": global_step}
|
|
swanlabcallback.on_log(args, state, control, logs=log)
|
|
swanlabcallback.on_train_end(args, state, control, model=model)
|
|
swanlabcallback._swanlab.finish()
|
|
os.environ.pop("SWANLAB_MODE", None)
|
|
shutil.rmtree(output_dir)
|
|
|
|
|
|
class TestTensorboardCallback(unittest.TestCase):
|
|
def test_tbcallback(self):
|
|
output_dir = tempfile.mkdtemp()
|
|
args = TrainingArguments(
|
|
output_dir=output_dir, max_steps=200, logging_steps=20, run_name="test_tbcallback", logging_dir=output_dir
|
|
)
|
|
state = TrainerState(trial_name="PaddleNLP")
|
|
control = TrainerControl()
|
|
tensorboard_callback = TensorBoardCallback()
|
|
self.assertIsNone(tensorboard_callback.tb_writer)
|
|
tensorboard_callback.on_train_begin(args, state, control)
|
|
try:
|
|
log_directory = tensorboard_callback.tb_writer.logdir
|
|
except AttributeError:
|
|
log_directory = tensorboard_callback.tb_writer.log_dir
|
|
self.assertEqual(log_directory, output_dir)
|
|
for global_step in range(args.max_steps):
|
|
state.global_step = global_step
|
|
if global_step % args.logging_steps == 0:
|
|
log = {"loss": 100 - 0.4 * global_step, "learning_rate": 0.1, "global_step": global_step}
|
|
tensorboard_callback.on_log(args, state, control, logs=log)
|
|
ea = event_accumulator.EventAccumulator(output_dir)
|
|
ea.Reload()
|
|
loss_scalars = ea.Scalars("train/loss")
|
|
learning_rate_scalars = ea.Scalars("train/learning_rate")
|
|
global_step_scalars = ea.Scalars("train/global_step")
|
|
for i, scalar in enumerate(loss_scalars):
|
|
expected_loss = 100 - 0.4 * scalar.step
|
|
self.assertAlmostEqual(scalar.value, expected_loss, places=5)
|
|
for i, scalar in enumerate(learning_rate_scalars):
|
|
expected_lr = 0.1
|
|
self.assertAlmostEqual(scalar.value, expected_lr, places=5)
|
|
for i, scalar in enumerate(global_step_scalars):
|
|
expected_step = i * args.logging_steps
|
|
self.assertEqual(scalar.value, expected_step)
|
|
tensorboard_callback.on_train_end(args, state, control)
|
|
self.assertIsNone(tensorboard_callback.tb_writer)
|
|
shutil.rmtree(output_dir)
|
|
|
|
|
|
class TestVisualDLCallback(unittest.TestCase):
|
|
def test_vdlcallback(self):
|
|
output_dir = tempfile.mkdtemp()
|
|
args = TrainingArguments(
|
|
output_dir=output_dir, max_steps=200, logging_steps=20, run_name="test_vdlcallback", logging_dir=output_dir
|
|
)
|
|
state = TrainerState(trial_name="PaddleNLP")
|
|
control = TrainerControl()
|
|
visualdl_callback = VisualDLCallback()
|
|
self.assertIsNone(visualdl_callback.vdl_writer)
|
|
visualdl_callback.on_train_begin(args, state, control)
|
|
self.assertEqual(visualdl_callback.vdl_writer.logdir, output_dir)
|
|
for global_step in range(args.max_steps):
|
|
state.global_step = global_step
|
|
if global_step % args.logging_steps == 0:
|
|
log = {"loss": 100 - 0.4 * global_step, "learning_rate": 0.1, "global_step": global_step}
|
|
visualdl_callback.on_log(args, state, control, logs=log)
|
|
reader = LogReader(file_path=visualdl_callback.vdl_writer.file_name)
|
|
loss_scalars = reader.get_data("scalar", "train/loss")
|
|
learning_rate_scalars = reader.get_data("scalar", "train/learning_rate")
|
|
global_step_scalars = reader.get_data("scalar", "train/global_step")
|
|
for i, scalar in enumerate(loss_scalars):
|
|
expected_loss = 100 - 0.4 * args.logging_steps * i
|
|
self.assertAlmostEqual(scalar.value, expected_loss, places=5)
|
|
for i, scalar in enumerate(learning_rate_scalars):
|
|
expected_lr = 0.1
|
|
self.assertAlmostEqual(scalar.value, expected_lr, places=5)
|
|
for i, scalar in enumerate(global_step_scalars):
|
|
expected_step = i * args.logging_steps
|
|
self.assertEqual(scalar.value, expected_step)
|
|
visualdl_callback.on_train_end(args, state, control)
|
|
self.assertIsNone(visualdl_callback.vdl_writer)
|
|
shutil.rmtree(output_dir)
|