trax
v1.4.1
Trax ist eine End-to-End-Bibliothek für Deep Learning, die sich auf klare Code und Geschwindigkeit konzentriert. Es wird aktiv verwendet und im Google Brain -Team gepflegt. Dieses Notebook (ausführen in Colab) zeigt, wie Sie Trax verwenden und wo Sie weitere Informationen finden können.
Wir begrüßen Beiträge zu Trax! Wir begrüßen PRS mit Code für neue Modelle und Ebenen sowie Verbesserungen unseres Codes und unserer Dokumentation. Wir lieben Notizbücher , die erklären, wie Modelle funktionieren, und zeigen, wie man sie verwendet, um Probleme zu lösen!
Hier sind einige Beispiel-Notizbücher:-
trax.data -APIAllgemeines Setup
Führen Sie die folgende Zelle (einmal) aus, bevor Sie eine der Code -Samples ausführen.
import os
import numpy as np
!p ip install - q - U trax
import trax So erstellen Sie einen englisch-deutschen Übersetzer in einigen Codezeilen:
# Create a Transformer model.
# Pre-trained model config in gs://trax-ml/models/translation/ende_wmt32k.gin
model = trax . models . Transformer (
input_vocab_size = 33300 ,
d_model = 512 , d_ff = 2048 ,
n_heads = 8 , n_encoder_layers = 6 , n_decoder_layers = 6 ,
max_len = 2048 , mode = 'predict' )
# Initialize using pre-trained weights.
model . init_from_file ( 'gs://trax-ml/models/translation/ende_wmt32k.pkl.gz' ,
weights_only = True )
# Tokenize a sentence.
sentence = 'It is nice to learn new things today!'
tokenized = list ( trax . data . tokenize ( iter ([ sentence ]), # Operates on streams.
vocab_dir = 'gs://trax-ml/vocabs/' ,
vocab_file = 'ende_32k.subword' ))[ 0 ]
# Decode from the Transformer.
tokenized = tokenized [ None , :] # Add batch dimension.
tokenized_translation = trax . supervised . decoding . autoregressive_sample (
model , tokenized , temperature = 0.0 ) # Higher temperature: more diverse results.
# De-tokenize,
tokenized_translation = tokenized_translation [ 0 ][: - 1 ] # Remove batch and EOS.
translation = trax . data . detokenize ( tokenized_translation ,
vocab_dir = 'gs://trax-ml/vocabs/' ,
vocab_file = 'ende_32k.subword' )
print ( translation ) Es ist schön, heute neue Dinge zu lernen!
TRAX enthält grundlegende Modelle (wie Resnet, LSTM, Transformator) und RL -Algorithmen (wie Stärke, A2C, PPO). Es wird auch aktiv für die Forschung verwendet und umfasst neue Modelle wie Reformer und neue RL -Algorithmen wie AWR. TRAX hat Bindungen an eine große Anzahl von Deep -Learning -Datensätzen, einschließlich Tensor2tensor- und Tensorflow -Datensätzen.
Sie können Trax entweder als Bibliothek aus Ihren eigenen Python -Skripten und Notizbüchern oder als Binäran der Shell verwenden, die für das Training großer Modelle bequemer sein kann. Es läuft ohne Änderungen an CPUs, GPUs und TPUs.
Sie können hier lernen, wie Trax funktioniert, wie Sie neue Modelle erstellen und wie Sie sie nach eigenen Daten trainieren können.
Die grundlegenden Einheiten, die durch Trax -Modelle fließen, sind Tensoren - mehrdimensionale Arrays, manchmal auch als Numpy -Arrays bezeichnet, aufgrund des am häufigsten verwendeten Pakets für Tensoroperationen - numpy . Sie sollten sich den Numpy -Leitfaden ansehen, wenn Sie nicht wissen, wie Sie mit Tensoren arbeiten: Trax verwendet auch die Numpy -API dafür.
In Trax möchten wir, dass Numpy -Operationen sehr schnell laufen und GPUs und TPUs verwenden, um sie zu beschleunigen. Wir möchten auch automatisch Funktionen von Funktionen auf Tensoren berechnen. Dies geschieht im trax.fastmath -Paket dank seiner Backends - Jax und TensorFlow Numpy.
from trax . fastmath import numpy as fastnp
trax . fastmath . use_backend ( 'jax' ) # Can be 'jax' or 'tensorflow-numpy'.
matrix = fastnp . array ([[ 1 , 2 , 3 ], [ 4 , 5 , 6 ], [ 7 , 8 , 9 ]])
print ( f'matrix = n { matrix } ' )
vector = fastnp . ones ( 3 )
print ( f'vector = { vector } ' )
product = fastnp . dot ( vector , matrix )
print ( f'product = { product } ' )
tanh = fastnp . tanh ( product )
print ( f'tanh(product) = { tanh } ' ) matrix =
[[1 2 3]
[4 5 6]
[7 8 9]]
vector = [1. 1. 1.]
product = [12. 15. 18.]
tanh(product) = [0.99999994 0.99999994 0.99999994]
Gradienten können mit trax.fastmath.grad berechnet werden.
def f ( x ):
return 2.0 * x * x
grad_f = trax . fastmath . grad ( f )
print ( f'grad(2x^2) at 1 = { grad_f ( 1.0 ) } ' ) grad(2x^2) at 1 = 4.0
Schichten sind grundlegende Bausteine von Trax -Modellen. Sie werden alles über sie in den Schichten intro lernen, aber im Moment sehen Sie sich die Implementierung einer Kern -Trax -Schicht an Embedding :
class Embedding ( base . Layer ):
"""Trainable layer that maps discrete tokens/IDs to vectors."""
def __init__ ( self ,
vocab_size ,
d_feature ,
kernel_initializer = init . RandomNormalInitializer ( 1.0 )):
"""Returns an embedding layer with given vocabulary size and vector size.
Args:
vocab_size: Size of the input vocabulary. The layer will assign a unique
vector to each ID in `range(vocab_size)`.
d_feature: Dimensionality/depth of the output vectors.
kernel_initializer: Function that creates (random) initial vectors for
the embedding.
"""
super (). __init__ ( name = f'Embedding_ { vocab_size } _ { d_feature } ' )
self . _d_feature = d_feature # feature dimensionality
self . _vocab_size = vocab_size
self . _kernel_initializer = kernel_initializer
def forward ( self , x ):
"""Returns embedding vectors corresponding to input token IDs.
Args:
x: Tensor of token IDs.
Returns:
Tensor of embedding vectors.
"""
return jnp . take ( self . weights , x , axis = 0 , mode = 'clip' )
def init_weights_and_state ( self , input_signature ):
"""Returns tensor of newly initialized embedding vectors."""
del input_signature
shape_w = ( self . _vocab_size , self . _d_feature )
w = self . _kernel_initializer ( shape_w , self . rng )
self . weights = w Schichten mit trainierbaren Gewichten wie Embedding müssen mit der Signatur (Form und DTYPE) des Eingangs initialisiert werden und können dann ausgeführt werden, indem sie aufgerufen werden.
from trax import layers as tl
# Create an input tensor x.
x = np . arange ( 15 )
print ( f'x = { x } ' )
# Create the embedding layer.
embedding = tl . Embedding ( vocab_size = 20 , d_feature = 32 )
embedding . init ( trax . shapes . signature ( x ))
# Run the layer -- y = embedding(x).
y = embedding ( x )
print ( f'shape of y = { y . shape } ' ) x = [ 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14]
shape of y = (15, 32)
Modelle in Trax werden aus Schichten hergestellt, die am häufigsten mit den Serial und Branch ausgeschaltet werden. Sie können mehr über diese Kombinators im Intro der Ebenen lesen und den Code für viele Modelle in trax/models/ , z. B. so implementiert, dass das Transformator -Sprachmodell implementiert wird. Im Folgenden finden Sie ein Beispiel dafür, wie Sie ein Stimmungsklassifizierungsmodell erstellen.
model = tl . Serial (
tl . Embedding ( vocab_size = 8192 , d_feature = 256 ),
tl . Mean ( axis = 1 ), # Average on axis 1 (length of sentence).
tl . Dense ( 2 ), # Classify 2 classes.
tl . LogSoftmax () # Produce log-probabilities.
)
# You can print model structure.
print ( model ) Serial[
Embedding_8192_256
Mean
Dense_2
LogSoftmax
]
Um Ihr Modell zu trainieren, benötigen Sie Daten. In Trax werden Datenströme als Python -Iteratoren dargestellt, sodass Sie next(data_stream) aufrufen und ein Tupel, z. B. (inputs, targets) erhalten. Mit TRAX können Sie TensorFlow -Datensätze problemlos verwenden und können auch einen Iterator aus Ihrer eigenen Textdatei mithilfe des Standards open('my_file.txt') .
train_stream = trax . data . TFDS ( 'imdb_reviews' , keys = ( 'text' , 'label' ), train = True )()
eval_stream = trax . data . TFDS ( 'imdb_reviews' , keys = ( 'text' , 'label' ), train = False )()
print ( next ( train_stream )) # See one example. (b"This was an absolutely terrible movie. Don't be lured in by Christopher Walken or Michael Ironside. Both are great actors, but this must simply be their worst role in history. Even their great acting could not redeem this movie's ridiculous storyline. This movie is an early nineties US propaganda piece. The most pathetic scenes were those when the Columbian rebels were making their cases for revolutions. Maria Conchita Alonso appeared phony, and her pseudo-love affair with Walken was nothing but a pathetic emotional plug in a movie that was devoid of any real meaning. I am disappointed that there are movies like this, ruining actor's like Christopher Walken's good name. I could barely sit through it.", 0)
Mit dem trax.data -Modul können Sie Input -Verarbeitungs -Pipelines, z. B. zum Tokenisieren und Mischen Ihrer Daten erstellen. Sie erstellen Datenpipelines mit trax.data.Serial und sind Funktionen, die Sie auf Streams anwenden, um verarbeitete Streams zu erstellen.
data_pipeline = trax . data . Serial (
trax . data . Tokenize ( vocab_file = 'en_8k.subword' , keys = [ 0 ]),
trax . data . Shuffle (),
trax . data . FilterByLength ( max_length = 2048 , length_keys = [ 0 ]),
trax . data . BucketByLength ( boundaries = [ 32 , 128 , 512 , 2048 ],
batch_sizes = [ 256 , 64 , 16 , 4 , 1 ],
length_keys = [ 0 ]),
trax . data . AddLossWeights ()
)
train_batches_stream = data_pipeline ( train_stream )
eval_batches_stream = data_pipeline ( eval_stream )
example_batch = next ( train_batches_stream )
print ( f'shapes = { [ x . shape for x in example_batch ] } ' ) # Check the shapes. shapes = [(4, 1024), (4,), (4,)]
Wenn Sie das Modell und die Daten haben, verwenden Sie trax.supervised.training , um Schulungs- und Bewertungsaufgaben zu definieren und eine Trainingsschleife zu erstellen. Die Trax -Trainingsschleife optimiert das Training und erstellt Tensorboard -Protokolle und Modellkontrollpunkte für Sie.
from trax . supervised import training
# Training task.
train_task = training . TrainTask (
labeled_data = train_batches_stream ,
loss_layer = tl . WeightedCategoryCrossEntropy (),
optimizer = trax . optimizers . Adam ( 0.01 ),
n_steps_per_checkpoint = 500 ,
)
# Evaluaton task.
eval_task = training . EvalTask (
labeled_data = eval_batches_stream ,
metrics = [ tl . WeightedCategoryCrossEntropy (), tl . WeightedCategoryAccuracy ()],
n_eval_batches = 20 # For less variance in eval numbers.
)
# Training loop saves checkpoints to output_dir.
output_dir = os . path . expanduser ( '~/output_dir/' )
!r m - rf { output_dir }
training_loop = training . Loop ( model ,
train_task ,
eval_tasks = [ eval_task ],
output_dir = output_dir )
# Run 2000 steps (batches).
training_loop . run ( 2000 ) Step 1: Ran 1 train steps in 0.78 secs
Step 1: train WeightedCategoryCrossEntropy | 1.33800304
Step 1: eval WeightedCategoryCrossEntropy | 0.71843582
Step 1: eval WeightedCategoryAccuracy | 0.56562500
Step 500: Ran 499 train steps in 5.77 secs
Step 500: train WeightedCategoryCrossEntropy | 0.62914723
Step 500: eval WeightedCategoryCrossEntropy | 0.49253047
Step 500: eval WeightedCategoryAccuracy | 0.74062500
Step 1000: Ran 500 train steps in 5.03 secs
Step 1000: train WeightedCategoryCrossEntropy | 0.42949259
Step 1000: eval WeightedCategoryCrossEntropy | 0.35451687
Step 1000: eval WeightedCategoryAccuracy | 0.83750000
Step 1500: Ran 500 train steps in 4.80 secs
Step 1500: train WeightedCategoryCrossEntropy | 0.41843575
Step 1500: eval WeightedCategoryCrossEntropy | 0.35207348
Step 1500: eval WeightedCategoryAccuracy | 0.82109375
Step 2000: Ran 500 train steps in 5.35 secs
Step 2000: train WeightedCategoryCrossEntropy | 0.38129005
Step 2000: eval WeightedCategoryCrossEntropy | 0.33760912
Step 2000: eval WeightedCategoryAccuracy | 0.85312500
Führen Sie nach dem Training das Modell wie jede Ebene aus, um Ergebnisse zu erzielen.
example_input = next ( eval_batches_stream )[ 0 ][ 0 ]
example_input_str = trax . data . detokenize ( example_input , vocab_file = 'en_8k.subword' )
print ( f'example input_str: { example_input_str } ' )
sentiment_log_probs = model ( example_input [ None , :]) # Add batch dimension.
print ( f'Model returned sentiment probabilities: { np . exp ( sentiment_log_probs ) } ' ) example input_str: I first saw this when I was a teen in my last year of Junior High. I was riveted to it! I loved the special effects, the fantastic places and the trial-aspect and flashback method of telling the story.<br /><br />Several years later I read the book and while it was interesting and I could definitely see what Swift was trying to say, I think that while it's not as perfect as the book for social commentary, as a story the movie is better. It makes more sense to have it be one long adventure than having Gulliver return after each voyage and making a profit by selling the tiny Lilliput sheep or whatever.<br /><br />It's much more arresting when everyone thinks he's crazy and the sheep DO make a cameo anyway. As a side note, when I saw Laputa I was stunned. It looks very much like the Kingdom of Zeal from the Chrono Trigger video game (1995) that also made me like this mini-series even more.<br /><br />I saw it again about 4 years ago, and realized that I still enjoyed it just as much. Really high quality stuff and began an excellent run of Sweeps mini-series for NBC who followed it up with the solid Merlin and interesting Alice in Wonderland.<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Model returned sentiment probabilities: [[3.984500e-04 9.996014e-01]]
