Spaces:
Configuration error
Configuration error
File size: 11,921 Bytes
a01ef8c |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 |
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (c) 2022 Intel Corporation
#
# 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.
#
# SPDX-License-Identifier: Apache-2.0
#
import os
import pytest
import shutil
import tempfile
from unittest.mock import MagicMock
from tlt.datasets import dataset_factory
from tlt.models import model_factory
try:
from tlt.datasets.text_classification.hf_custom_text_classification_dataset import HFCustomTextClassificationDataset
except ModuleNotFoundError:
print("WARNING: Unable to import HFCustomTextClassificationDataset.")
@pytest.mark.integration
@pytest.mark.pytorch
@pytest.mark.parametrize('model_name,dataset_name,extra_layers,correct_num_layers,test_inc',
[['bert-base-cased', 'imdb', None, 1, False],
['distilbert-base-uncased', 'imdb', [384, 192], 5, True]])
def test_pyt_text_classification(model_name, dataset_name, extra_layers, correct_num_layers, test_inc):
"""
Tests basic transfer learning functionality for PyTorch text classification models using a hugging face dataset
"""
framework = 'pytorch'
output_dir = tempfile.mkdtemp()
# Get the dataset
dataset = dataset_factory.get_dataset(output_dir, 'text_classification', framework, dataset_name,
'huggingface', split=["train"], shuffle_files=False)
# Get the model
model = model_factory.get_model(model_name, framework)
# Preprocess the dataset
dataset.preprocess(model_name, batch_size=32)
dataset.shuffle_split(train_pct=0.02, val_pct=0.01, seed=6)
assert dataset._validation_type == 'shuffle_split'
# Evaluate before training
pretrained_metrics = model.evaluate(dataset)
assert len(pretrained_metrics) > 0
# Train
train_history = model.train(dataset, output_dir=output_dir, epochs=1, do_eval=False, extra_layers=extra_layers)
assert train_history is not None and isinstance(train_history, dict)
assert 'Loss' in train_history
assert 'Acc' in train_history
assert 'train_runtime' in train_history
assert 'train_samples_per_second' in train_history
classifier_layer = getattr(model._model, "classifier")
try:
# If extra_layers given, the classifier is a Sequential layer with given input
n_layers = len(classifier_layer)
except TypeError:
# If not given, the classifer is just a single Linear layer
n_layers = 1
assert n_layers == correct_num_layers
# Evaluate
trained_metrics = model.evaluate(dataset)
assert trained_metrics['eval_loss'] <= pretrained_metrics['eval_loss']
assert trained_metrics['eval_accuracy'] >= pretrained_metrics['eval_accuracy']
# Export the saved model
saved_model_dir = model.export(output_dir)
assert os.path.isdir(saved_model_dir)
assert os.path.isfile(os.path.join(saved_model_dir, "model.pt"))
# Reload the saved model
reload_model = model_factory.get_model(model_name, framework)
reload_model.load_from_directory(saved_model_dir)
# Evaluate
reload_metrics = reload_model.evaluate(dataset)
assert reload_metrics['eval_accuracy'] == trained_metrics['eval_accuracy']
# Ensure we get 'NotImplementedError' for graph_optimization
with pytest.raises(NotImplementedError):
model.optimize_graph(os.path.join(saved_model_dir, 'optimized'))
# Quantization
if test_inc:
inc_output_dir = os.path.join(output_dir, "quantized", model_name)
os.makedirs(inc_output_dir, exist_ok=True)
model.quantize(inc_output_dir, dataset)
assert os.path.exists(os.path.join(inc_output_dir, "model.pt"))
# Delete the temp output directory
if os.path.exists(output_dir) and os.path.isdir(output_dir):
shutil.rmtree(output_dir)
@pytest.mark.integration
@pytest.mark.pytorch
@pytest.mark.parametrize('model_name',
['bert-base-cased'])
def test_custom_dataset_workflow(model_name):
"""
Tests the full workflow for PYT text classification using a custom dataset mock
"""
model = model_factory.get_model(model_name, framework='pytorch', use_case="text_classification")
output_dir = tempfile.mkdtemp()
os.environ["TORCH_HOME"] = output_dir
mock_dataset = MagicMock()
mock_dataset.__class__ = HFCustomTextClassificationDataset
mock_dataset.validation_subset = ['fun', 'terrible']
mock_dataset.train_subset = ["fun, happy, boring, terrible"]
mock_dataset.class_names = ['good', 'bad']
# Preprocess the dataset and split to get small subsets for training and validation
mock_dataset.shuffle_split(train_pct=0.1, val_pct=0.1, shuffle_files=False)
mock_dataset.preprocess(model_name, batch_size=32)
# Train for 1 epoch
history = model.train(mock_dataset, output_dir=output_dir, epochs=1, seed=10, do_eval=False)
assert history is not None
# Evaluate
model.evaluate(mock_dataset)
# export the saved model
saved_model_dir = model.export(output_dir)
assert os.path.isdir(saved_model_dir)
assert os.path.isfile(os.path.join(saved_model_dir, "model.pt"))
# Reload the saved model
reload_model = model_factory.get_model(model_name, 'pytorch')
reload_model.load_from_directory(saved_model_dir)
# Evaluate
metrics = reload_model.evaluate(mock_dataset)
assert len(metrics) > 0
# Delete the temp output directory
if os.path.exists(output_dir) and os.path.isdir(output_dir):
shutil.rmtree(output_dir)
@pytest.mark.integration
@pytest.mark.pytorch
@pytest.mark.parametrize('model_name,dataset_name',
[['distilbert-base-uncased', 'imdb']])
def test_initial_checkpoints(model_name, dataset_name):
framework = 'pytorch'
output_dir = tempfile.mkdtemp()
checkpoint_dir = os.path.join(output_dir, model_name + '_checkpoints')
# Get the dataset
dataset = dataset_factory.get_dataset(output_dir, 'text_classification', framework, dataset_name,
'huggingface', split=["train"], shuffle_files=False)
# Get the model
model = model_factory.get_model(model_name, framework)
assert model._generate_checkpoints is True
dataset.preprocess(model_name, batch_size=32)
dataset.shuffle_split(train_pct=0.01, val_pct=0.01, seed=10)
# Train
model.train(dataset, output_dir=output_dir, epochs=2, do_eval=False)
trained_metrics = model.evaluate(dataset)
# Delete the model and train a brand new model but instead we resume training from checkpoints
del model
model = model_factory.get_model(model_name, framework)
model.train(dataset, output_dir=output_dir, epochs=2, do_eval=False,
initial_checkpoints=os.path.join(checkpoint_dir, 'checkpoint.pt'))
improved_metrics = model.evaluate(dataset)
assert improved_metrics['eval_loss'] < trained_metrics['eval_loss']
assert improved_metrics['eval_accuracy'] > trained_metrics['eval_accuracy']
# Delete the temp output directory
if os.path.exists(output_dir) and os.path.isdir(output_dir):
shutil.rmtree(output_dir)
@pytest.mark.integration
@pytest.mark.pytorch
@pytest.mark.parametrize('model_name,dataset_name',
[['distilbert-base-uncased', 'imdb']])
def test_freeze_bert(model_name, dataset_name):
framework = 'pytorch'
output_dir = tempfile.mkdtemp()
# Get the dataset
dataset = dataset_factory.get_dataset(output_dir, 'text_classification', framework, dataset_name,
'huggingface', split=["train"], shuffle_files=False)
# Get the model
model = model_factory.get_model(model_name, framework)
dataset.preprocess(model_name, batch_size=32)
dataset.shuffle_split(train_pct=0.01, val_pct=0.01, seed=10)
# Train
model.train(dataset, output_dir=output_dir, epochs=1, do_eval=False)
# Freeze feature layers
layer_name = "features"
model.freeze_layer(layer_name)
# Check everything is frozen (not trainable) in the layer
for (name, module) in model._model.named_children():
if name == layer_name:
for param in module.parameters():
assert param.requires_grad is False
# Delete the temp output directory
if os.path.exists(output_dir) and os.path.isdir(output_dir):
shutil.rmtree(output_dir)
@pytest.mark.integration
@pytest.mark.pytorch
@pytest.mark.parametrize('model_name,dataset_name',
[['distilbert-base-uncased', 'imdb']])
def test_unfreeze_bert(model_name, dataset_name):
framework = 'pytorch'
output_dir = tempfile.mkdtemp()
# Get the dataset
dataset = dataset_factory.get_dataset(output_dir, 'text_classification', framework, dataset_name,
'huggingface', split=["train"], shuffle_files=False)
# Get the model
model = model_factory.get_model(model_name, framework)
dataset.preprocess(model_name, batch_size=32)
dataset.shuffle_split(train_pct=0.01, val_pct=0.01, seed=10)
# Train
model.train(dataset, output_dir=output_dir, epochs=1, do_eval=False)
layer_name = "features"
model.unfreeze_layer(layer_name)
# Check everything is unfrozen (trainable) in the layer
for (name, module) in model._model.named_children():
if name == layer_name:
for param in module.parameters():
assert param.requires_grad is True
# Delete the temp output directory
if os.path.exists(output_dir) and os.path.isdir(output_dir):
shutil.rmtree(output_dir)
@pytest.mark.integration
@pytest.mark.pytorch
@pytest.mark.parametrize('model_name,dataset_name',
[['distilbert-base-uncased', 'imdb']])
def test_list_layers_bert(model_name, dataset_name):
import io
import unittest.mock as mock
framework = 'pytorch'
output_dir = tempfile.mkdtemp()
# Get the model
model = model_factory.get_model(model_name, framework)
# Get the dataset
dataset = dataset_factory.get_dataset(output_dir, 'text_classification', framework, dataset_name,
'huggingface', split=["train"], shuffle_files=False)
dataset.preprocess(model_name, batch_size=32)
dataset.shuffle_split(train_pct=0.01, val_pct=0.01, seed=10)
# Train
model.train(dataset, output_dir=output_dir, epochs=1, do_eval=False)
# Mock stdout and sterr to capture the function's output
stdout = io.StringIO()
stderr = io.StringIO()
with mock.patch('sys.stdout', stdout), mock.patch('sys.stderr', stderr):
model.list_layers(verbose=True)
# Assert the function printed the correct output of the trainable layers
output = stdout.getvalue().strip()
assert 'distilbert' in output
assert 'embeddings: 23835648/23835648 parameters are trainable' in output
assert 'transformer: 42527232/42527232 parameters are trainable' in output
assert 'pre_classifier: 590592/590592 parameters are trainable' in output
assert 'dropout: 0/0 parameters are trainable' in output
assert 'dropout: 0/0 parameters are trainable' in output
assert 'Total Trainable Parameters: 66955010/66955010' in output
# Delete the temp output directory
if os.path.exists(output_dir) and os.path.isdir(output_dir):
shutil.rmtree(output_dir)
|