trax
v1.4.1
Trax เป็นไลบรารีแบบ end-to-end สำหรับการเรียนรู้ลึกที่มุ่งเน้นไปที่รหัสและความเร็วที่ชัดเจน มันถูกใช้อย่างแข็งขันและบำรุงรักษาในทีม Google Brain สมุดบันทึกนี้ (เรียกใช้ใน colab) แสดงวิธีใช้ TRAX และสถานที่ที่คุณสามารถค้นหาข้อมูลเพิ่มเติมได้
เรายินดีต้อนรับ การมีส่วนร่วม กับ Trax! เรายินดีต้อนรับ PRS ด้วยรหัสสำหรับรุ่นใหม่และเลเยอร์รวมถึงการปรับปรุงรหัสและเอกสารของเรา เรารัก สมุดบันทึก ที่อธิบายว่าโมเดลทำงานอย่างไรและแสดงวิธีการใช้เพื่อแก้ปัญหา!
นี่คือสมุดบันทึกตัวอย่างบางส่วน:-
trax.data apiการตั้งค่าทั่วไป
ดำเนินการเซลล์ต่อไปนี้ (หนึ่งครั้ง) ก่อนที่จะเรียกใช้ตัวอย่างรหัสใด ๆ
import os
import numpy as np
!p ip install - q - U trax
import trax นี่คือวิธีที่คุณสร้างนักแปลภาษาอังกฤษ-เยอรมันในรหัสไม่กี่บรรทัด:
# 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 มีโมเดลพื้นฐาน (เช่น Resnet, LSTM, Transformer) และอัลกอริทึม RL (เช่น Senreforce, A2C, PPO) นอกจากนี้ยังใช้อย่างแข็งขันสำหรับการวิจัยและรวมถึงโมเดลใหม่เช่นนักปฏิรูปและอัลกอริทึม RL ใหม่เช่น AWR Trax มีการเชื่อมโยงชุดข้อมูลการเรียนรู้ลึกจำนวนมากรวมถึงชุดข้อมูล Tensor2Tensor และ TensorFlow
คุณสามารถใช้ TRAX ไม่ว่าจะเป็นไลบรารีจากสคริปต์ Python และสมุดบันทึกของคุณเองหรือเป็นไบนารีจากเปลือกซึ่งสะดวกกว่าสำหรับการฝึกอบรมรุ่นใหญ่ มันทำงานโดยไม่มีการเปลี่ยนแปลงใด ๆ ในซีพียู GPU และ TPU
คุณสามารถเรียนรู้ที่นี่ว่า Trax ทำงานอย่างไรวิธีสร้างโมเดลใหม่และวิธีการฝึกอบรมด้วยข้อมูลของคุณเอง
หน่วยพื้นฐานที่ไหลผ่านโมเดล TRAX คือ เทนเซอร์ - อาร์เรย์หลายมิติซึ่งบางครั้งก็เรียกว่าอาร์เรย์ NUMPY เนื่องจากแพ็คเกจที่ใช้กันอย่างแพร่หลายที่สุดสำหรับการทำงานของเทนเซอร์ - numpy คุณควรดูคำแนะนำ NumPy หากคุณไม่ทราบวิธีการทำงานบนเทนเซอร์: Trax ยังใช้ NumPy API สำหรับสิ่งนั้น
ใน TRAX เราต้องการให้การดำเนินการ NumPy ทำงานได้อย่างรวดเร็วการใช้ GPU และ TPU เพื่อเร่งความเร็ว นอกจากนี้เรายังต้องการคำนวณการไล่ระดับสีของฟังก์ชั่นบนเทนเซอร์โดยอัตโนมัติ สิ่งนี้ทำในแพ็คเกจ trax.fastmath ด้วยแบ็กเอนด์ - Jax และ 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]
การไล่ระดับสีสามารถคำนวณได้โดยใช้ trax.fastmath.grad
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
เลเยอร์เป็นหน่วยการสร้างพื้นฐานของโมเดล TRAX คุณจะได้เรียนรู้ทั้งหมดเกี่ยวกับพวกเขาในอินโทรเลเยอร์ แต่ตอนนี้เพียงแค่ดูการใช้งานของเลเยอร์ Trax Core หนึ่งชั้น 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 เลเยอร์ที่มีน้ำหนักฝึกอบรมได้เช่น Embedding จะต้องเริ่มต้นด้วยลายเซ็น (รูปร่างและ dtype) ของอินพุตจากนั้นสามารถเรียกใช้ได้โดยเรียกพวกเขา
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)
แบบจำลองใน TRAX ถูกสร้างขึ้นจากเลเยอร์ส่วนใหญ่มักใช้เครื่องผสม Serial และ Branch คุณสามารถอ่านเพิ่มเติมเกี่ยวกับ combinators เหล่านั้นในชั้นนำและดูรหัสสำหรับหลายรุ่นใน trax/models/ , EG นี่คือวิธีที่โมเดลภาษาหม้อแปลงถูกนำไปใช้ ด้านล่างเป็นตัวอย่างของวิธีการสร้างรูปแบบการจำแนกความเชื่อมั่น
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
]
ในการฝึกอบรมแบบจำลองของคุณคุณต้องมีข้อมูล ใน TRAX สตรีมข้อมูลจะถูกแสดงเป็นตัววนซ้ำ Python ดังนั้นคุณสามารถโทรหา next(data_stream) และรับ tuple, เช่น (inputs, targets) Trax ช่วยให้คุณใช้ชุดข้อมูล TensorFlow ได้อย่างง่ายดายและคุณยังสามารถรับตัววนซ้ำจากไฟล์ข้อความของคุณเองโดยใช้มาตรฐาน 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)
การใช้โมดูล trax.data คุณสามารถสร้างท่อประมวลผลอินพุตเช่น tokenize และสลับข้อมูลของคุณ คุณสร้างท่อข้อมูลโดยใช้ trax.data.Serial และเป็นฟังก์ชั่นที่คุณใช้กับสตรีมเพื่อสร้างสตรีมที่ประมวลผล
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,)]
เมื่อคุณมีโมเดลและข้อมูลให้ใช้ trax.supervised.training เพื่อกำหนดงานการฝึกอบรมและการประเมินและสร้างห่วงการฝึกอบรม TRAX Training Loop เพิ่มประสิทธิภาพการฝึกอบรมและจะสร้างบันทึก Tensorboard และจุดตรวจสอบรุ่นสำหรับคุณ
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
หลังจากฝึกอบรมแบบจำลองให้เรียกใช้เหมือนเลเยอร์ใด ๆ เพื่อให้ได้ผลลัพธ์
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]]
