Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
transformers
Commits
20d6931e
Unverified
Commit
20d6931e
authored
Apr 30, 2021
by
Matt
Committed by
GitHub
Apr 30, 2021
Browse files
Update TF text classification example (#11496)
Big refactor, fixes and multi-GPU/TPU support
parent
8b945ef0
Changes
3
Expand all
Show whitespace changes
Inline
Side-by-side
Showing
3 changed files
with
201 additions
and
187 deletions
+201
-187
examples/tensorflow/text-classification/README.md
examples/tensorflow/text-classification/README.md
+14
-0
examples/tensorflow/text-classification/run_text_classification.py
...tensorflow/text-classification/run_text_classification.py
+182
-185
src/transformers/training_args_tf.py
src/transformers/training_args_tf.py
+5
-2
No files found.
examples/tensorflow/text-classification/README.md
View file @
20d6931e
...
@@ -54,6 +54,20 @@ After training, the model will be saved to `--output_dir`. Once your model is tr
...
@@ -54,6 +54,20 @@ After training, the model will be saved to `--output_dir`. Once your model is tr
by calling the script without a
`--train_file`
or
`--validation_file`
; simply pass it the output_dir containing
by calling the script without a
`--train_file`
or
`--validation_file`
; simply pass it the output_dir containing
the trained model and a
`--test_file`
and it will write its predictions to a text file for you.
the trained model and a
`--test_file`
and it will write its predictions to a text file for you.
### Multi-GPU and TPU usage
By default, the script uses a
`MirroredStrategy`
and will use multiple GPUs effectively if they are available. TPUs
can also be used by passing the name of the TPU resource with the
`--tpu`
argument.
### Memory usage and data loading
One thing to note is that all data is loaded into memory in this script. Most text classification datasets are small
enough that this is not an issue, but if you have a very large dataset you will need to modify the script to handle
data streaming. This is particularly challenging for TPUs, given the stricter requirements and the sheer volume of data
required to keep them fed. A full explanation of all the possible pitfalls is a bit beyond this example script and
README, but for more information you can see the 'Input Datasets' section of
[
this document
](
https://www.tensorflow.org/guide/tpu
)
.
### Example command
### Example command
```
```
python run_text_classification.py \
python run_text_classification.py \
...
...
examples/tensorflow/text-classification/run_text_classification.py
View file @
20d6931e
This diff is collapsed.
Click to expand it.
src/transformers/training_args_tf.py
View file @
20d6931e
...
@@ -212,6 +212,9 @@ class TFTrainingArguments(TrainingArguments):
...
@@ -212,6 +212,9 @@ class TFTrainingArguments(TrainingArguments):
else
:
else
:
tpu
=
tf
.
distribute
.
cluster_resolver
.
TPUClusterResolver
()
tpu
=
tf
.
distribute
.
cluster_resolver
.
TPUClusterResolver
()
except
ValueError
:
except
ValueError
:
if
self
.
tpu_name
:
raise
RuntimeError
(
f
"Couldn't connect to TPU
{
self
.
tpu_name
}
!"
)
else
:
tpu
=
None
tpu
=
None
if
tpu
:
if
tpu
:
...
@@ -233,7 +236,7 @@ class TFTrainingArguments(TrainingArguments):
...
@@ -233,7 +236,7 @@ class TFTrainingArguments(TrainingArguments):
# If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
# If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
strategy
=
tf
.
distribute
.
MirroredStrategy
()
strategy
=
tf
.
distribute
.
MirroredStrategy
()
else
:
else
:
raise
ValueError
(
"Cannot find the proper strategy please check your environment properties."
)
raise
ValueError
(
"Cannot find the proper strategy
,
please check your environment properties."
)
return
strategy
return
strategy
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment