import trax
from trax import layers as tl
from trax.supervised import training
# Trax offers the WideResnet architecture in it's models module
from trax.models.resnet import WideResnet
%%capture
train_stream = trax.data.TFDS('cifar10', keys=('image', 'label'), train=True)()
eval_stream = trax.data.TFDS('cifar10', keys=('image', 'label'), train=False)()
train_data_pipeline = trax.data.Serial(
trax.data.Shuffle(),
trax.data.Batch(64),
trax.data.AddLossWeights(),
)
train_batches_stream = train_data_pipeline(train_stream)
eval_data_pipeline = trax.data.Serial(
trax.data.Batch(64),
trax.data.AddLossWeights(),
)
eval_batches_stream = eval_data_pipeline(eval_stream)
thin_model = tl.Serial(
WideResnet(widen_factor = 1),
tl.LogSoftmax()
)
wide_model = tl.Serial(
WideResnet(widen_factor = 2),
tl.LogSoftmax()
)
wider_model = tl.Serial(
WideResnet(widen_factor = 3),
tl.LogSoftmax()
)
widest_model = tl.Serial(
WideResnet(widen_factor = 4),
tl.LogSoftmax()
)
train_task = training.TrainTask(
labeled_data=train_batches_stream,
loss_layer=tl.CrossEntropyLoss(),
optimizer=trax.optimizers.Adam(0.01),
n_steps_per_checkpoint=1000,
)
eval_task = training.EvalTask(
labeled_data=eval_batches_stream,
metrics=[tl.CrossEntropyLoss(), tl.Accuracy()],
n_eval_batches=20,
)
training_loop = training.Loop(thin_model,
train_task,
eval_tasks=[eval_task],
output_dir='./thin_model')
training_loop.run(5000)
Step 1: Ran 1 train steps in 50.80 secs
Step 1: train CrossEntropyLoss | 2.23222899
Step 1: eval CrossEntropyLoss | 2.29001374
Step 1: eval Accuracy | 0.17578125
Step 1000: Ran 999 train steps in 450.67 secs
Step 1000: train CrossEntropyLoss | 1.51696503
Step 1000: eval CrossEntropyLoss | 1.24389263
Step 1000: eval Accuracy | 0.55156250
Step 2000: Ran 1000 train steps in 391.61 secs
Step 2000: train CrossEntropyLoss | 1.14578664
Step 2000: eval CrossEntropyLoss | 1.07358952
Step 2000: eval Accuracy | 0.60625000
Step 3000: Ran 1000 train steps in 379.49 secs
Step 3000: train CrossEntropyLoss | 0.98207539
Step 3000: eval CrossEntropyLoss | 0.97521102
Step 3000: eval Accuracy | 0.66406250
Step 4000: Ran 1000 train steps in 375.78 secs
Step 4000: train CrossEntropyLoss | 0.87311232
Step 4000: eval CrossEntropyLoss | 0.86088786
Step 4000: eval Accuracy | 0.71328125
Step 5000: Ran 1000 train steps in 382.36 secs
Step 5000: train CrossEntropyLoss | 0.79590148
Step 5000: eval CrossEntropyLoss | 0.90162431
Step 5000: eval Accuracy | 0.67968750
training_loop = training.Loop(wide_model,
train_task,
eval_tasks=[eval_task],
output_dir='./wide_model')
training_loop.run(5000)
Step 1: Ran 1 train steps in 47.65 secs
Step 1: train CrossEntropyLoss | 2.45481586
Step 1: eval CrossEntropyLoss | 2.29830315
Step 1: eval Accuracy | 0.17265625
Step 1000: Ran 999 train steps in 912.74 secs
Step 1000: train CrossEntropyLoss | 1.53430188
Step 1000: eval CrossEntropyLoss | 1.26175007
Step 1000: eval Accuracy | 0.54296875
Step 2000: Ran 1000 train steps in 898.69 secs
Step 2000: train CrossEntropyLoss | 1.13339007
Step 2000: eval CrossEntropyLoss | 0.98637019
Step 2000: eval Accuracy | 0.65390625
Step 3000: Ran 1000 train steps in 914.89 secs
Step 3000: train CrossEntropyLoss | 0.94035697
Step 3000: eval CrossEntropyLoss | 0.89243511
Step 3000: eval Accuracy | 0.68906250
Step 4000: Ran 1000 train steps in 878.78 secs
Step 4000: train CrossEntropyLoss | 0.82797289
Step 4000: eval CrossEntropyLoss | 0.86628049
Step 4000: eval Accuracy | 0.70156250
Step 5000: Ran 1000 train steps in 896.67 secs
Step 5000: train CrossEntropyLoss | 0.74435514
Step 5000: eval CrossEntropyLoss | 0.82093265
Step 5000: eval Accuracy | 0.70859375
training_loop = training.Loop(wider_model,
train_task,
eval_tasks=[eval_task],
output_dir='./wider_model')
training_loop.run(5000)
Step 1: Ran 1 train steps in 49.64 secs
Step 1: train CrossEntropyLoss | 2.46369076
Step 1: eval CrossEntropyLoss | 2.42145765
Step 1: eval Accuracy | 0.16015625
Step 1000: Ran 999 train steps in 1462.86 secs
Step 1000: train CrossEntropyLoss | 1.55000281
Step 1000: eval CrossEntropyLoss | 1.31225752
Step 1000: eval Accuracy | 0.53203125
Step 2000: Ran 1000 train steps in 1417.44 secs
Step 2000: train CrossEntropyLoss | 1.14296257
Step 2000: eval CrossEntropyLoss | 1.05580651
Step 2000: eval Accuracy | 0.61796875
Step 3000: Ran 1000 train steps in 1412.55 secs
Step 3000: train CrossEntropyLoss | 0.96064937
Step 3000: eval CrossEntropyLoss | 0.91904441
Step 3000: eval Accuracy | 0.66093750
Step 4000: Ran 1000 train steps in 1394.65 secs
Step 4000: train CrossEntropyLoss | 0.86051035
Step 4000: eval CrossEntropyLoss | 0.79895681
Step 4000: eval Accuracy | 0.71875000
Step 5000: Ran 1000 train steps in 1325.14 secs
Step 5000: train CrossEntropyLoss | 0.76998872
Step 5000: eval CrossEntropyLoss | 0.79824924
Step 5000: eval Accuracy | 0.73046875
training_loop = training.Loop(widest_model,
train_task,
eval_tasks=[eval_task],
output_dir='./widest_model')
training_loop.run(5000)
Step 1: Ran 1 train steps in 50.33 secs
Step 1: train CrossEntropyLoss | 2.46660376
Step 1: eval CrossEntropyLoss | 2.54380413
Step 1: eval Accuracy | 0.18828125
Step 1000: Ran 999 train steps in 2232.45 secs
Step 1000: train CrossEntropyLoss | 1.60101640
Step 1000: eval CrossEntropyLoss | 1.32047499
Step 1000: eval Accuracy | 0.51562500
Step 2000: Ran 1000 train steps in 2207.08 secs
Step 2000: train CrossEntropyLoss | 1.23905230
Step 2000: eval CrossEntropyLoss | 1.14217659
Step 2000: eval Accuracy | 0.59140625
Step 3000: Ran 1000 train steps in 2193.96 secs
Step 3000: train CrossEntropyLoss | 1.02972627
Step 3000: eval CrossEntropyLoss | 0.95948886
Step 3000: eval Accuracy | 0.66015625
Step 4000: Ran 1000 train steps in 2184.98 secs
Step 4000: train CrossEntropyLoss | 0.88306051
Step 4000: eval CrossEntropyLoss | 0.92772388
Step 4000: eval Accuracy | 0.66796875
Step 5000: Ran 1000 train steps in 2179.92 secs
Step 5000: train CrossEntropyLoss | 0.79552704
Step 5000: eval CrossEntropyLoss | 0.73757971
Step 5000: eval Accuracy | 0.72968750