Las traducciones son generadas a través de traducción automática. En caso de conflicto entre la traducción y la version original de inglés, prevalecerá la version en inglés.
# Usar un modelo de aprendizaje profundo propio
**importante**
Amazon Web Services (AWS) anuncia que no habrá nuevas versiones o versiones de SageMaker Training Compiler. Puede seguir utilizando SageMaker Training Compiler a través de los AWS Deep Learning Containers (DLCs) existentes para SageMaker formación. Es importante tener en cuenta que, si bien los existentes DLCs permanecen accesibles, ya no recibirán parches ni actualizaciones de ellos AWS, de acuerdo con la [Política de soporte de AWS Deep Learning Containers Framework](https://docs.aws.amazon.com/deep-learning-containers/latest/devguide/support-policy.html).
En esta guía se explica cómo adaptar el script de entrenamiento para un trabajo de entrenamiento acelerado por compilador. La preparación del script de entrenamiento depende de lo siguiente:
+ Entornos de entrenamiento, como la entrenamiento distribuida o de un solo kernel.
+ Marcos y bibliotecas que se utilizan para crear el script de entrenamiento.
Elija uno de los siguientes temas en función del marco de trabajo que utilice.
**Topics**
+ [PyTorch](training-compiler-pytorch-models.md)
+ [TensorFlow](training-compiler-tensorflow.md)
**nota**
Cuando termine de preparar su guion de formación, podrá realizar un trabajo de SageMaker formación utilizando las clases de cálculo del marco de SageMaker IA. Para obtener más información, consulte el tema anterior en [Habilita el compilador SageMaker de entrenamiento](training-compiler-enable.md).
# PyTorch
Incorpore su propio PyTorch modelo a la SageMaker IA y ejecute el trabajo de formación con SageMaker Training Compiler.
**Topics**
+ [PyTorch Modelos con Hugging Face Transformers](#training-compiler-pytorch-models-transformers)
## PyTorch Modelos con Hugging Face Transformers
PyTorch [los modelos con [Hugging Face](https://huggingface.co/docs/transformers/index) Transformers se PyTorch basan en la API torch.NN.Module.](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) Hugging Face Transformers también [ofrece](https://huggingface.co/docs/transformers/main_classes/trainer) clases de modelos para instructores y previamente entrenados PyTorch para ayudar a reducir el esfuerzo de configuración de los modelos de procesamiento del lenguaje natural (NLP). Tras preparar el guion de formación, podrá iniciar un trabajo de formación con la SageMaker IA `PyTorch` o con un `HuggingFace` estimador con la configuración del compilador de SageMaker formación. Después, pasará al siguiente tema en. [Habilita el compilador SageMaker de entrenamiento](training-compiler-enable.md)
**sugerencia**
Cuando cree un tokenizador para un modelo de PNL con Transformers en su script de entrenamiento, asegúrese de utilizar una forma de tensor de entrada estática especificando `padding='max_length'`. No utilice `padding='longest'` porque rellenar la secuencia más larga del lote puede cambiar la forma del tensor de cada lote de entrenamiento. La forma de entrada dinámica puede desencadenar la recompilación del modelo y aumentar el tiempo total de entrenamiento. Para obtener más información sobre las opciones de relleno de los tokenizadores de Transformers, consulte [Padding and truncation](https://huggingface.co/docs/transformers/pad_truncation) en la *documentación de Hugging Face Transformers*.
**Topics**
+ [Grandes modelos lingüísticos mediante la clase `Trainer` de Hugging Face Transformers](#training-compiler-pytorch-models-transformers-trainer)
+ [Modelos de idiomas de gran tamaño que se utilizan PyTorch directamente (sin la API Hugging Face Transformers Trainer)](#training-compiler-pytorch-models-non-trainer)
### Grandes modelos lingüísticos mediante la clase `Trainer` de Hugging Face Transformers
Si utilizas la clase de entrenador de la biblioteca de Transformers, no necesitas realizar ningún cambio adicional en tu guion de formación. SageMaker Training Compiler compila automáticamente tu modelo Trainer si lo habilitas en la clase Estimador. El siguiente código muestra la forma básica de un guion de PyTorch entrenamiento con la API Hugging Face Trainer.
```
from transformers import Trainer, TrainingArguments
training_args=TrainingArguments(**kwargs)
trainer=Trainer(args=training_args, **kwargs)
```
**Topics**
+ [Para el entrenamiento con una sola GPU](#training-compiler-pytorch-models-transformers-trainer-single-gpu)
+ [Para entrenamiento distribuido](#training-compiler-pytorch-models-transformers-trainer-distributed)
+ [Mejores prácticas para utilizar Training Compiler con SageMaker `Trainer`](#training-compiler-pytorch-models-transformers-trainer-best-practices)
#### Para el entrenamiento con una sola GPU
No necesita cambiar su código cuando utilice la clase [https://huggingface.co/docs/transformers/main_classes/trainer](https://huggingface.co/docs/transformers/main_classes/trainer).
#### Para entrenamiento distribuido
**PyTorch v1.11.0 y versiones posteriores**
Para ejecutar un entrenamiento distribuido con SageMaker Training Compiler, debe agregar la siguiente `_mp_fn()` función a su script de entrenamiento y empaquetar la función. `main()` Redirige las llamadas a `_mp_fn(index)` funciones del motor de ejecución distribuido por la SageMaker IA para PyTorch (`pytorchxla`) a la `main()` función de tu guion de entrenamiento.
```
def _mp_fn(index):
main()
```
Esta función acepta el argumento `index` para indicar el rango de la GPU actual en el clúster para el entrenamiento distribuido. Para encontrar más scripts de ejemplo, consulte los scripts de ejemplo de modelado lingüístico de [Hugging Face Transformers](https://github.com/huggingface/transformers/blob/v4.21.1/examples/pytorch/language-modeling).
**Para Transformers v4.17 y anteriores con v1.10.2 y anteriores PyTorch **
SageMaker Training Compiler utiliza un mecanismo alternativo para lanzar un trabajo de formación distribuido y no es necesario realizar ninguna modificación en el guion de formación. En cambio, SageMaker Training Compiler requiere que pases un guion de un lanzador de entrenamiento distribuido por SageMaker IA al `entry_point` argumento y que pases tu guion de entrenamiento al `hyperparameters` argumento del estimador Hugging Face de SageMaker IA.
#### Mejores prácticas para utilizar Training Compiler con SageMaker `Trainer`
+ [Asegúrese de utilizar SyncFree optimizadores estableciendo el `optim` argumento en al configurar los `adamw_torch_xla` transformadores. TrainingArgument](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments). Consulte también [Optimizer](https://huggingface.co/docs/transformers/v4.23.1/en/perf_train_gpu_one#optimizer) en la *documentación de Hugging Face Transformers*.
+ Asegúrese de que el rendimiento del canal de procesamiento de datos es superior al rendimiento de entrenamiento. Puedes modificar los `preprocessing_num_workers` argumentos `dataloader_num_workers` y de los [transformadores. TrainingArgument](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments)clase para lograr esto. Por lo general, estos deben ser mayores o iguales que el número de GPUs pero menores que el número de CPUs.
Cuando haya terminado de adaptar su script de entrenamiento, continúe con [Ejecute trabajos PyTorch de entrenamiento con SageMaker Training Compiler](training-compiler-enable-pytorch.md).
### Modelos de idiomas de gran tamaño que se utilizan PyTorch directamente (sin la API Hugging Face Transformers Trainer)
Si tiene un script de entrenamiento que se usa PyTorch directamente, debe realizar cambios adicionales en el script de PyTorch entrenamiento para implementar /XLA. PyTorch Siga las instrucciones para modificar el script y configurar correctamente las primitivas de /XLA. PyTorch
**Topics**
+ [Para el entrenamiento con una sola GPU](#training-compiler-pytorch-models-non-trainer-single-gpu)
+ [Para entrenamiento distribuido](#training-compiler-pytorch-models-non-trainer-distributed)
+ [Mejores prácticas para usar Training Compiler con /XLA SageMaker PyTorch](#training-compiler-pytorch-models-best-practices)
#### Para el entrenamiento con una sola GPU
1. Importe las bibliotecas de optimización.
```
import torch_xla
import torch_xla.core.xla_model as xm
```
1. Cambiar el dispositivo de destino para que sea XLA en lugar de `torch.device("cuda")`
```
device=xm.xla_device()
```
1. Si utiliza PyTorch la [precisión mixta automática](https://pytorch.org/docs/stable/amp.html) (AMP), haga lo siguiente:
1. Reemplace `torch.cuda.amp` por lo siguiente:
```
import torch_xla.amp
```
1. Reemplace `torch.optim.SGD` y `torch.optim.Adam` por lo siguiente:
```
import torch_xla.amp.syncfree.Adam as adam
import torch_xla.amp.syncfree.SGD as SGD
```
1. Reemplace `torch.cuda.amp.GradScaler` por lo siguiente:
```
import torch_xla.amp.GradScaler as grad_scaler
```
1. Si no utiliza AMP, reemplace `optimizer.step()` por lo siguiente:
```
xm.optimizer_step(optimizer)
```
1. Si utilizas un cargador de datos distribuido, incluye tu cargador de datos en la clase /XLA PyTorch: `ParallelLoader`
```
import torch_xla.distributed.parallel_loader as pl
parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)
```
1. Añada `mark_step` al final del bucle de entrenamiento cuando no estés utilizando `parallel_loader`:
```
xm.mark_step()
```
1. Para comprobar tu entrenamiento, utiliza el método de puntos de control del modelo /XLA: PyTorch
```
xm.save(model.state_dict(), path_to_save)
```
Cuando haya terminado de adaptar su script de entrenamiento, continúe con [Ejecute trabajos PyTorch de entrenamiento con SageMaker Training Compiler](training-compiler-enable-pytorch.md).
#### Para entrenamiento distribuido
Además de los cambios enumerados en la [Para el entrenamiento con una sola GPU](#training-compiler-pytorch-models-non-trainer-single-gpu) sección anterior, añada los siguientes cambios para distribuir correctamente la carga de trabajo entre todos. GPUs
1. Si esta utilizando AMP, añada `all_reduce` después de `scaler.scale(loss).backward()`:
```
gradients=xm._fetch_gradients(optimizer)
xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())
```
1. Si necesita establecer variables para `local_ranks` y `world_size`, utilice un código similar al siguiente:
```
local_rank=xm.get_local_ordinal()
world_size=xm.xrt_world_size()
```
1. Para cualquier `world_size` (`num_gpus_per_node*num_nodes`) mayor que `1`, debe definir un muestreador de entrenamiento que debe tener un aspecto similar al siguiente:
```
import torch_xla.core.xla_model as xm
if xm.xrt_world_size() > 1:
train_sampler=torch.utils.data.distributed.DistributedSampler(
train_dataset,
num_replicas=xm.xrt_world_size(),
rank=xm.get_ordinal(),
shuffle=True
)
train_loader=torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
sampler=train_sampler,
drop_last=args.drop_last,
shuffle=False if train_sampler else True,
num_workers=args.num_workers
)
```
1. Realice los siguientes cambios para asegurarse de que utiliza el `parallel_loader` proporcionado por el módulo `torch_xla distributed`.
```
import torch_xla.distributed.parallel_loader as pl
train_device_loader=pl.MpDeviceLoader(train_loader, device)
```
`train_device_loader`Funciona como un PyTorch cargador normal de la siguiente manera:
```
for step, (data, target) in enumerate(train_device_loader):
optimizer.zero_grad()
output=model(data)
loss=torch.nn.NLLLoss(output, target)
loss.backward()
```
Con todos estos cambios, deberías poder lanzar un entrenamiento distribuido con cualquier PyTorch modelo sin la API de Transformer Trainer. Tenga en cuenta que estas instrucciones se pueden utilizar tanto para varias GPU de un solo nodo como para varias GPU de varios nodos.
1. **Para la PyTorch versión 1.11.0 y versiones posteriores**
Para ejecutar un entrenamiento distribuido con SageMaker Training Compiler, debe agregar la siguiente `_mp_fn()` función a su script de entrenamiento y empaquetar la función. `main()` Redirige las llamadas a `_mp_fn(index)` funciones del motor de ejecución distribuido por la SageMaker IA para PyTorch (`pytorchxla`) a la `main()` función de tu guion de entrenamiento.
```
def _mp_fn(index):
main()
```
Esta función acepta el argumento `index` para indicar el rango de la GPU actual en el clúster para el entrenamiento distribuido. Para encontrar más scripts de ejemplo, consulte los scripts de ejemplo de modelado lingüístico de [Hugging Face Transformers](https://github.com/huggingface/transformers/blob/v4.21.1/examples/pytorch/language-modeling).
**Para Transformers v4.17 y anteriores con v1.10.2 y anteriores PyTorch**
SageMaker Training Compiler utiliza un mecanismo alternativo para lanzar un trabajo de formación distribuido y requiere que pases un guion de un lanzador de entrenamiento distribuido de SageMaker IA al `entry_point` argumento y que pases tu guion de entrenamiento al `hyperparameters` argumento del estimador Hugging Face de SageMaker IA.
Cuando haya terminado de adaptar su script de entrenamiento, continúe con [Ejecute trabajos PyTorch de entrenamiento con SageMaker Training Compiler](training-compiler-enable-pytorch.md).
#### Mejores prácticas para usar Training Compiler con /XLA SageMaker PyTorch
[Si desea utilizar el compilador de SageMaker formación en su script de PyTorch formación nativo, puede que primero desee familiarizarse con PyTorch el uso de dispositivos XLA.](https://pytorch.org/xla/release/1.9/index.html) En las siguientes secciones se enumeran algunas de las mejores prácticas para habilitar XLA. PyTorch
**nota**
En esta sección de prácticas recomendadas se parte del supuesto de que se utilizan los siguientes PyTorch/XLA módulos:
```
import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
```
##### Comprenda el modo perezoso en PyTorch /XLA
Una diferencia significativa entre el sistema nativo PyTorch/XLA y el sistema nativo PyTorch es que el PyTorch/XLA sistema se ejecuta en modo lento, mientras que el sistema nativo PyTorch se ejecuta en modo ansioso. Los tensores en modo Lazy son marcadores de posición para construir el gráfico computacional hasta que se materializan después de que la compilación y la evaluación se hayan completado. El PyTorch/XLA sistema crea el gráfico computacional sobre la marcha cuando usted llama PyTorch APIs para crear el cálculo utilizando tensores y operadores. El gráfico computacional se compila y ejecuta cuando `xm.mark_step()` es llamado explícita o implícitamente por `pl.MpDeviceLoader/pl.ParallelLoader`, o cuando se solicita explícitamente el valor de un tensor como por ejemplo llamando a `loss.item()` o `print(loss)`.
##### Minimice la cantidad de usos y *compilation-and-executions*`pl.MpDeviceLoader/pl.ParallelLoader``xm.step_closure`
Para obtener el mejor rendimiento, debe tener en cuenta las posibles formas de iniciar, tal *compilation-and-executions*como se describe en, [Comprenda el modo perezoso en PyTorch /XLA](#training-compiler-pytorch-models-best-practices-lazy-mode) y debe intentar minimizar la cantidad de compilation-and-executions. Lo ideal es que solo compilation-and-execution sea necesaria una por iteración de entrenamiento y que se inicie automáticamente por`pl.MpDeviceLoader/pl.ParallelLoader`. El `MpDeviceLoader` está optimizado para XLA y debe utilizarse siempre que sea posible para obtener el mejor rendimiento. Durante el entrenamiento, es posible que desee examinar algunos resultados intermedios, como los valores de pérdida. En tal caso, la impresión de tensores perezosos debe envolverse con tensores perezosos `xm.add_step_closure()` para evitar que sea innecesaria. compilation-and-executions
##### Utilizar AMP y optimizadores `syncfree`
Entrenar en el modo Automatic Mixed Precision (AMP) acelera considerablemente tu velocidad de entrenamiento al aprovechar los núcleos Tensor de NVIDIA. GPUs SageMaker Training Compiler proporciona `syncfree` optimizadores optimizados para XLA a fin de mejorar el rendimiento de AMP. Actualmente, están disponibles los tres optimizadores `syncfree` siguientes, que deben utilizarse si es posible para obtener el mejor rendimiento.
```
torch_xla.amp.syncfree.SGD
torch_xla.amp.syncfree.Adam
torch_xla.amp.syncfree.AdamW
```
Estos optimizadores `syncfree` deben emparejarse con `torch_xla.amp.GradScaler` para el escalado/desescalado de gradientes.
**sugerencia**
A partir de la PyTorch versión 1.13.1, SageMaker Training Compiler mejora el rendimiento PyTorch/XLA al permitir anular automáticamente los optimizadores (como SGD, Adam, AdamW) incluidos en las versiones sin sincronización (como,,) `torch.optim` o `transformers.optimization` con las versiones sin sincronización de los mismos. `torch_xla.amp.syncfree` `torch_xla.amp.syncfree.SGD` `torch_xla.amp.syncfree.Adam` `torch_xla.amp.syncfree.AdamW` No tiene que cambiar las líneas de código en las que define los optimizadores en su script de entrenamiento.
# TensorFlow
Incorpore su propio TensorFlow modelo a la SageMaker IA y ejecute el trabajo de formación con SageMaker Training Compiler.
## TensorFlow Modelos
SageMaker Training Compiler optimiza automáticamente las cargas de trabajo de entrenamiento de modelos que se crean sobre la TensorFlow API nativa o la API de Keras de alto nivel.
**sugerencia**
Para preprocesar el conjunto de datos de entrada, asegúrese de utilizar una forma de entrada estática. La forma de entrada dinámica puede iniciar la recompilación del modelo y aumentar el tiempo total de entrenamiento.
### Uso de Keras (recomendado)
[Para obtener la mejor aceleración del compilador, recomendamos utilizar modelos que sean subclases de Keras (tf.Keras.model). TensorFlow ](https://www.tensorflow.org/api_docs/python/tf/keras/Model)
#### Para el entrenamiento con una sola GPU
No es necesario realizar ningún cambio adicional en el script de entrenamiento.
### Sin Keras
SageMaker Training Compiler no admite la ejecución intensiva en. TensorFlow En consecuencia, deberías combinar tu modelo y tus ciclos de entrenamiento con la TensorFlow función decorator (`@tf.function`) para aprovechar la aceleración del compilador.
SageMaker [Training Compiler realiza una optimización a nivel de gráfico y utiliza el decorador para asegurarse de que las TensorFlow funciones estén configuradas para ejecutarse en modo gráfico.](https://www.tensorflow.org/guide/intro_to_graphs)
#### Para el entrenamiento con una sola GPU
TensorFlow La versión 2.0 o posterior tiene activada la ejecución rápida de forma predeterminada, por lo que debes añadir el `@tf.function` decorador delante de cada función que utilices para construir un modelo. TensorFlow
## TensorFlow Modelos con Hugging Face Transformers
TensorFlow los modelos con [Hugging Face](https://huggingface.co/docs/transformers/index) Transformers se TensorFlow basan en la API [tf.Keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model). Hugging Face Transformers también ofrece clases de modelos previamente entrenadas TensorFlow para ayudar a reducir el esfuerzo de configurar los modelos de procesamiento del lenguaje natural (NLP). Tras crear su propio guion de entrenamiento con la biblioteca de Transformers, puede ejecutar el guion de entrenamiento mediante el `HuggingFace` estimador de SageMaker IA con la clase de configuración SageMaker Training Compiler, tal y como se muestra en el tema anterior en. [Ejecute trabajos TensorFlow de entrenamiento con SageMaker Training Compiler](training-compiler-enable-tensorflow.md)
SageMaker Training Compiler optimiza automáticamente las cargas de trabajo de entrenamiento de modelos que se basan en la TensorFlow API nativa o en la API de alto nivel de Keras, como los modelos de transformadores. TensorFlow
**sugerencia**
Cuando cree un tokenizador para un modelo de PNL con Transformers en su script de entrenamiento, asegúrese de utilizar una forma de tensor de entrada estática especificando `padding='max_length'`. No utilice `padding='longest'` porque rellenar la secuencia más larga del lote puede cambiar la forma del tensor de cada lote de entrenamiento. La forma de entrada dinámica puede iniciar la recompilación del modelo y aumentar el tiempo total de entrenamiento. Para obtener más información sobre las opciones de relleno de los tokenizadores de Transformers, consulte [Padding and truncation](https://huggingface.co/docs/transformers/pad_truncation) en la *documentación de Hugging Face Transformers*.
**Topics**
+ [Uso de Keras](#training-compiler-tensorflow-models-transformers-keras)
+ [Sin Keras](#training-compiler-tensorflow-models-transformers-no-keras)
### Uso de Keras
[Para obtener la mejor aceleración del compilador, recomendamos utilizar modelos que sean subclases de Keras (tf.Keras.model). TensorFlow ](https://www.tensorflow.org/api_docs/python/tf/keras/Model) Como se indica en la página de [recorrido rápido](https://huggingface.co/docs/transformers/quicktour) de la documentación de *Hugging Face Transformers*, puede utilizar los modelos como modelos Keras TensorFlow normales.
#### Para el entrenamiento con una sola GPU
No es necesario realizar ningún cambio adicional en el script de entrenamiento.
#### Para entrenamiento distribuido
SageMaker La aceleración del compilador de entrenamiento funciona de forma transparente para cargas de trabajo con varias GPU cuando el modelo se construye y entrena con Keras dentro del alcance de la demanda. APIs [https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy)
1. Elija la estrategia de entrenamiento distribuida adecuada.
1. En el caso de varias GPU de un solo nodo, utilíice `tf.distribute.MirroredStrategy` para establecer la estrategia.
```
strategy = tf.distribute.MirroredStrategy()
```
1. En el caso de varios nodos y varias GPU, añade el siguiente código para configurar correctamente la configuración de entrenamiento TensorFlow distribuido antes de crear la estrategia.
```
def set_sm_dist_config():
DEFAULT_PORT = '8890'
DEFAULT_CONFIG_FILE = '/opt/ml/input/config/resourceconfig.json'
with open(DEFAULT_CONFIG_FILE) as f:
config = json.loads(f.read())
current_host = config['current_host']
tf_config = {
'cluster': {
'worker': []
},
'task': {'type': 'worker', 'index': -1}
}
for i, host in enumerate(config['hosts']):
tf_config['cluster']['worker'].append("%s:%s" % (host, DEFAULT_PORT))
if current_host == host:
tf_config['task']['index'] = i
os.environ['TF_CONFIG'] = json.dumps(tf_config)
set_sm_dist_config()
```
Utilice `tf.distribute.MultiWorkerMirroredStrategy` para establecer la estrategia.
```
strategy = tf.distribute.MultiWorkerMirroredStrategy()
```
1. Usando la estrategia que elija, encapsule el modelo.
```
with strategy.scope():
# create a model and do fit
```
### Sin Keras
Si quieres crear modelos personalizados con bucles de entrenamiento personalizados TensorFlow sin usar Keras, deberías combinar el modelo y el ciclo de entrenamiento con la TensorFlow función decorator (`@tf.function`) para aprovechar la aceleración del compilador.
SageMaker Training Compiler realiza una optimización a nivel de gráfico y utiliza el decorador para asegurarse de que las TensorFlow funciones estén configuradas para ejecutarse en modo gráfico.
#### Para el entrenamiento con una sola GPU
TensorFlow La versión 2.0 o posterior tiene activada la ejecución rápida de forma predeterminada, por lo que debes añadir el `@tf.function` decorador delante de cada función que utilices para construir un modelo. TensorFlow
#### Para entrenamiento distribuido
Además de los cambios necesarios para [usar Keras para el entrenamiento distribuido](https://docs.aws.amazon.com/sagemaker/latest/dg/training-compiler-tensorflow-models.html#training-compiler-tensorflow-models-transformers-keras), debe asegurarse de que las funciones que se ejecutarán en cada GPU estén anotadas con `@tf.function`, mientras que las funciones de comunicación entre GPU no estén anotadas. El ejemplo deñ código de entrenamiento debería ser como el siguiente:
```
@tf.function()
def compiled_step(inputs, outputs):
with tf.GradientTape() as tape:
pred=model(inputs, training=True)
total_loss=loss_object(outputs, pred)/args.batch_size
gradients=tape.gradient(total_loss, model.trainable_variables)
return total_loss, pred, gradients
def train_step(inputs, outputs):
total_loss, pred, gradients=compiled_step(inputs, outputs)
if args.weight_decay > 0.:
gradients=[g+v*args.weight_decay for g,v in zip(gradients, model.trainable_variables)]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss.update_state(total_loss)
train_accuracy.update_state(outputs, pred)
@tf.function()
def train_step_dist(inputs, outputs):
strategy.run(train_step, args= (inputs, outputs))
```
Tenga en cuenta que esta instrucción se puede utilizar tanto para varias GPU de un solo nodo como para varias GPU de varios nodos.