[add] model

config.json

@@ -0,0 +1,332 @@
+{
+  "_commit_hash": null,
+  "_name_or_path": "facebook/sam-vit-huge",
+  "architectures": [
+    "ScaMultitaskV2Model"
+  ],
+  "cache_dir": "/mnt/blob/weights/.model.cache/",
+  "initializer_range": 0.02,
+  "mask_caption_decoder_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "additional_num_hidden_layers": 12,
+    "architectures": null,
+    "attention_downsample_rate": 2,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "relu",
+    "hidden_size": 256,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "iou_head_depth": 3,
+    "iou_head_hidden_dim": 256,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "mlp_dim": 2048,
+    "model_type": "",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 8,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_caption_heads": 1,
+    "num_caption_tokens": 8,
+    "num_hidden_layers": 2,
+    "num_multimask_outputs": 3,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "model_type": "sca",
+  "num_task_tokens": 6,
+  "prompt_encoder_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_size": 256,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_embedding_size": 64,
+    "image_size": 1024,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "mask_input_channels": 16,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_point_embeddings": 4,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 16,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "text_config": {
+    "_name_or_path": "openlm-research/open_llama_3b_v2",
+    "add_cross_attention": false,
+    "architectures": [
+      "LlamaForCausalLM"
+    ],
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "silu",
+    "hidden_size": 3200,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 8640,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 2048,
+    "min_length": 0,
+    "model_type": "llama",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 32,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 26,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-06,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "float16",
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 32000
+  },
+  "tie_word_embeddings": false,
+  "torch_dtype": "float16",
+  "transformers_version": null,
+  "use_decoder_only_language_model": true,
+  "vision_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "global_attn_indexes": [
+      7,
+      15,
+      23,
+      31
+    ],
+    "hidden_act": "gelu",
+    "hidden_size": 1280,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 1024,
+    "initializer_factor": 1.0,
+    "initializer_range": 1e-10,
+    "intermediate_size": 6144,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "mlp_dim": 5120,
+    "mlp_ratio": 4.0,
+    "model_type": "",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 32,
+    "num_pos_feats": 128,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_channels": 256,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 16,
+    "prefix": null,
+    "problem_type": null,
+    "projection_dim": 512,
+    "pruned_heads": {},
+    "qkv_bias": true,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_abs_pos": true,
+    "use_bfloat16": false,
+    "use_rel_pos": true,
+    "window_size": 14
+  }
+}

latest

@@ -0,0 +1 @@
+global_step200000

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf13798c876ccaa5889358468b0595fed1fa74febd499115f701c9bad80b9365
+size 8175399446

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "</s>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91b289e85fa20fd375d8b33dc12f77616f18abc6359804471d1fafcb425fecb8
+size 511574

tokenizer_config.json

@@ -0,0 +1,33 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "model_max_length": 20,
+  "pad_token": null,
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

trainer_state.json

@@ -0,0 +1,2194 @@
+{
+  "best_metric": 2.9787235260009766,
+  "best_model_checkpoint": "/mnt/output/projects/sca-xiaoke-v3/amlt-results/7300886584.15971-e315970f-15b9-410a-b0a9-3912402cdf8b/checkpoint-195000",
+  "epoch": 82.67879288962381,
+  "global_step": 200000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "_prepare_inputs_in_ms": 14.42616805434227,
+      "compute_loss_in_ms": 1299.8907640576363,
+      "epoch": 0.0,
+      "learning_rate/full": 0.0,
+      "loss": 9.016,
+      "step": 1,
+      "training_step_in_ms": 2032.6302126049995
+    },
+    {
+      "epoch": 0.0,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 9.1282958984375,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.03355821582361717,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 189.4348,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 4.223,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.132,
+      "step": 1
+    },
+    {
+      "_prepare_inputs_in_ms": 4.084589669218985,
+      "compute_loss_in_ms": 284.03653898784466,
+      "epoch": 0.41,
+      "learning_rate/full": 0.00039998495845181817,
+      "loss": 3.8908,
+      "step": 1000,
+      "training_step_in_ms": 1082.7972516678535
+    },
+    {
+      "_prepare_inputs_in_ms": 4.054305288940668,
+      "compute_loss_in_ms": 284.15025370568037,
+      "epoch": 0.83,
+      "learning_rate/full": 0.0003999216713877652,
+      "loss": 3.4805,
+      "step": 2000,
+      "training_step_in_ms": 1086.274579320103
+    },
+    {
+      "_prepare_inputs_in_ms": 4.071794345974922,
+      "compute_loss_in_ms": 284.45942908525467,
+      "epoch": 1.24,
+      "learning_rate/full": 0.00039980895784128267,
+      "loss": 3.405,
+      "step": 3000,
+      "training_step_in_ms": 1119.1077427528799
+    },
+    {
+      "_prepare_inputs_in_ms": 4.085798408836126,
+      "compute_loss_in_ms": 284.8285736106336,
+      "epoch": 1.65,
+      "learning_rate/full": 0.00039964684567845476,
+      "loss": 3.3537,
+      "step": 4000,
+      "training_step_in_ms": 1150.6125138737261
+    },
+    {
+      "_prepare_inputs_in_ms": 4.0938322730362415,
+      "compute_loss_in_ms": 285.10985093563795,
+      "epoch": 2.07,
+      "learning_rate/full": 0.0003994358471466495,
+      "loss": 3.3218,
+      "step": 5000,
+      "training_step_in_ms": 1108.4499053694308
+    },
+    {
+      "epoch": 2.07,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.269989252090454,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.21715476097311398,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 111.5411,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.172,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.224,
+      "step": 5000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.128877357011888,
+      "compute_loss_in_ms": 285.1207774169743,
+      "epoch": 2.48,
+      "learning_rate/full": 0.0003991751687428334,
+      "loss": 3.2918,
+      "step": 6000,
+      "training_step_in_ms": 1120.3654654994607
+    },
+    {
+      "_prepare_inputs_in_ms": 4.121768821030855,
+      "compute_loss_in_ms": 284.9130438826978,
+      "epoch": 2.89,
+      "learning_rate/full": 0.0003988659173490642,
+      "loss": 3.2675,
+      "step": 7000,
+      "training_step_in_ms": 1130.791154742241
+    },
+    {
+      "_prepare_inputs_in_ms": 4.129346951842308,
+      "compute_loss_in_ms": 285.35304405912757,
+      "epoch": 3.31,
+      "learning_rate/full": 0.0003985069299623724,
+      "loss": 3.2451,
+      "step": 8000,
+      "training_step_in_ms": 1133.7338739708066
+    },
+    {
+      "_prepare_inputs_in_ms": 4.133303381502628,
+      "compute_loss_in_ms": 285.2391963750124,
+      "epoch": 3.72,
+      "learning_rate/full": 0.0003980992984040504,
+      "loss": 3.2334,
+      "step": 9000,
+      "training_step_in_ms": 1086.6298492662609
+    },
+    {
+      "_prepare_inputs_in_ms": 4.132662046700716,
+      "compute_loss_in_ms": 285.1939390525222,
+      "epoch": 4.13,
+      "learning_rate/full": 0.00039764230739017226,
+      "loss": 3.2122,
+      "step": 10000,
+      "training_step_in_ms": 1111.2217365466058
+    },
+    {
+      "epoch": 4.13,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.1718945503234863,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.23451229725961228,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.8619,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.349,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 10000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.151588522079514,
+      "compute_loss_in_ms": 285.39055866748095,
+      "epoch": 4.55,
+      "learning_rate/full": 0.00039713751381134497,
+      "loss": 3.1979,
+      "step": 11000,
+      "training_step_in_ms": 1089.2837468609214
+    },
+    {
+      "_prepare_inputs_in_ms": 4.139887526631355,
+      "compute_loss_in_ms": 285.3860225379467,
+      "epoch": 4.96,
+      "learning_rate/full": 0.000396583019288311,
+      "loss": 3.1933,
+      "step": 12000,
+      "training_step_in_ms": 1084.1714271605015
+    },
+    {
+      "_prepare_inputs_in_ms": 4.176147662103176,
+      "compute_loss_in_ms": 285.6273371577263,
+      "epoch": 5.37,
+      "learning_rate/full": 0.0003959805510184613,
+      "loss": 3.1765,
+      "step": 13000,
+      "training_step_in_ms": 1089.3080548346043
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1766685508191586,
+      "compute_loss_in_ms": 285.3720509596169,
+      "epoch": 5.79,
+      "learning_rate/full": 0.00039532972748016767,
+      "loss": 3.1727,
+      "step": 14000,
+      "training_step_in_ms": 1091.7199603579938
+    },
+    {
+      "_prepare_inputs_in_ms": 4.191178072243929,
+      "compute_loss_in_ms": 285.4580160602927,
+      "epoch": 6.2,
+      "learning_rate/full": 0.0003946307092543998,
+      "loss": 3.1591,
+      "step": 15000,
+      "training_step_in_ms": 1090.6033144891262
+    },
+    {
+      "epoch": 6.2,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.1219239234924316,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.24624792238728155,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.8497,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.35,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 15000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.180258305334464,
+      "compute_loss_in_ms": 285.4601420760155,
+      "epoch": 6.61,
+      "learning_rate/full": 0.0003938828970266217,
+      "loss": 3.1562,
+      "step": 16000,
+      "training_step_in_ms": 1084.842809855938
+    },
+    {
+      "_prepare_inputs_in_ms": 4.194100107997656,
+      "compute_loss_in_ms": 285.48905945569277,
+      "epoch": 7.03,
+      "learning_rate/full": 0.00039308797090204444,
+      "loss": 3.1508,
+      "step": 17000,
+      "training_step_in_ms": 1085.5392471551895
+    },
+    {
+      "_prepare_inputs_in_ms": 4.201232250779867,
+      "compute_loss_in_ms": 285.61364733427763,
+      "epoch": 7.44,
+      "learning_rate/full": 0.0003922445359987763,
+      "loss": 3.1333,
+      "step": 18000,
+      "training_step_in_ms": 1091.4973263852298
+    },
+    {
+      "_prepare_inputs_in_ms": 4.189023811370134,
+      "compute_loss_in_ms": 285.4624082148075,
+      "epoch": 7.85,
+      "learning_rate/full": 0.0003913554018411121,
+      "loss": 3.1351,
+      "step": 19000,
+      "training_step_in_ms": 1095.9623138792813
+    },
+    {
+      "_prepare_inputs_in_ms": 4.215472485870123,
+      "compute_loss_in_ms": 285.6230415776372,
+      "epoch": 8.27,
+      "learning_rate/full": 0.00039041818639024787,
+      "loss": 3.1297,
+      "step": 20000,
+      "training_step_in_ms": 1185.9847482070327
+    },
+    {
+      "epoch": 8.27,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.094672918319702,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.24975866124736495,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.4566,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.376,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.231,
+      "step": 20000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1831618874538234,
+      "compute_loss_in_ms": 285.3183429725468,
+      "epoch": 8.68,
+      "learning_rate/full": 0.00038943398810118026,
+      "loss": 3.119,
+      "step": 21000,
+      "training_step_in_ms": 1082.7558356113732
+    },
+    {
+      "_prepare_inputs_in_ms": 4.174146838486195,
+      "compute_loss_in_ms": 285.4058397859335,
+      "epoch": 9.09,
+      "learning_rate/full": 0.0003884019945070803,
+      "loss": 3.1178,
+      "step": 22000,
+      "training_step_in_ms": 1087.0350129008293
+    },
+    {
+      "_prepare_inputs_in_ms": 4.204949229955673,
+      "compute_loss_in_ms": 285.3504670076072,
+      "epoch": 9.51,
+      "learning_rate/full": 0.00038732452418171673,
+      "loss": 3.1093,
+      "step": 23000,
+      "training_step_in_ms": 1090.7997342124581
+    },
+    {
+      "_prepare_inputs_in_ms": 4.196985870599747,
+      "compute_loss_in_ms": 285.4016271494329,
+      "epoch": 9.92,
+      "learning_rate/full": 0.0003862008343330083,
+      "loss": 3.1062,
+      "step": 24000,
+      "training_step_in_ms": 1090.5269000642002
+    },
+    {
+      "_prepare_inputs_in_ms": 4.206760194152594,
+      "compute_loss_in_ms": 285.6653628349304,
+      "epoch": 10.33,
+      "learning_rate/full": 0.00038503000849003844,
+      "loss": 3.0972,
+      "step": 25000,
+      "training_step_in_ms": 1090.913136728108
+    },
+    {
+      "epoch": 10.33,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.075335741043091,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.25707208633373096,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.472,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.375,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 25000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.204354642367944,
+      "compute_loss_in_ms": 285.44070146232843,
+      "epoch": 10.75,
+      "learning_rate/full": 0.0003838159164157488,
+      "loss": 3.097,
+      "step": 26000,
+      "training_step_in_ms": 1085.0929874032736
+    },
+    {
+      "_prepare_inputs_in_ms": 4.202797457575798,
+      "compute_loss_in_ms": 285.55905482545495,
+      "epoch": 11.16,
+      "learning_rate/full": 0.00038255527679000744,
+      "loss": 3.0903,
+      "step": 27000,
+      "training_step_in_ms": 1089.4926370121539
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2012519761919975,
+      "compute_loss_in_ms": 285.71758703514934,
+      "epoch": 11.58,
+      "learning_rate/full": 0.0003812482649321827,
+      "loss": 3.0892,
+      "step": 28000,
+      "training_step_in_ms": 1090.198759533465
+    },
+    {
+      "_prepare_inputs_in_ms": 4.220829274505377,
+      "compute_loss_in_ms": 285.5986107811332,
+      "epoch": 11.99,
+      "learning_rate/full": 0.0003798978172979138,
+      "loss": 3.0889,
+      "step": 29000,
+      "training_step_in_ms": 1085.0686310827732
+    },
+    {
+      "_prepare_inputs_in_ms": 4.217640113085508,
+      "compute_loss_in_ms": 285.6399230584502,
+      "epoch": 12.4,
+      "learning_rate/full": 0.0003785029825690954,
+      "loss": 3.0749,
+      "step": 30000,
+      "training_step_in_ms": 1087.9900991134346
+    },
+    {
+      "epoch": 12.4,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.058861017227173,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.25600252799309026,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.6245,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.298,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 30000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.224213841484814,
+      "compute_loss_in_ms": 285.3009058833122,
+      "epoch": 12.82,
+      "learning_rate/full": 0.00037706410490032555,
+      "loss": 3.0794,
+      "step": 31000,
+      "training_step_in_ms": 1087.378763064742
+    },
+    {
+      "_prepare_inputs_in_ms": 4.233129996806383,
+      "compute_loss_in_ms": 285.4745088033378,
+      "epoch": 13.23,
+      "learning_rate/full": 0.0003755815393131386,
+      "loss": 3.074,
+      "step": 32000,
+      "training_step_in_ms": 1091.7205754183233
+    },
+    {
+      "_prepare_inputs_in_ms": 4.223707340657711,
+      "compute_loss_in_ms": 285.1553194858134,
+      "epoch": 13.64,
+      "learning_rate/full": 0.0003740556516084091,
+      "loss": 3.0686,
+      "step": 33000,
+      "training_step_in_ms": 1087.8120190612972
+    },
+    {
+      "_prepare_inputs_in_ms": 4.245889626443386,
+      "compute_loss_in_ms": 285.67459550127387,
+      "epoch": 14.06,
+      "learning_rate/full": 0.0003724852264850082,
+      "loss": 3.0733,
+      "step": 34000,
+      "training_step_in_ms": 1086.3751963675022
+    },
+    {
+      "_prepare_inputs_in_ms": 4.233293745666742,
+      "compute_loss_in_ms": 285.4981838874519,
+      "epoch": 14.47,
+      "learning_rate/full": 0.00037087542640234865,
+      "loss": 3.064,
+      "step": 35000,
+      "training_step_in_ms": 1089.8129360377789
+    },
+    {
+      "epoch": 14.47,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.052946090698242,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2584196718918565,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.4779,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.375,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 35000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.194417025257901,
+      "compute_loss_in_ms": 285.25339871644974,
+      "epoch": 14.88,
+      "learning_rate/full": 0.00036922019737873653,
+      "loss": 3.064,
+      "step": 36000,
+      "training_step_in_ms": 1084.3368335030973
+    },
+    {
+      "_prepare_inputs_in_ms": 4.225800335407257,
+      "compute_loss_in_ms": 285.52921985834837,
+      "epoch": 15.3,
+      "learning_rate/full": 0.00036752484999829976,
+      "loss": 3.0581,
+      "step": 37000,
+      "training_step_in_ms": 1090.5179475583136
+    },
+    {
+      "_prepare_inputs_in_ms": 4.208773214370012,
+      "compute_loss_in_ms": 285.2426546551287,
+      "epoch": 15.71,
+      "learning_rate/full": 0.0003657881683678541,
+      "loss": 3.0581,
+      "step": 38000,
+      "training_step_in_ms": 1089.0830878019333
+    },
+    {
+      "_prepare_inputs_in_ms": 4.219740275293589,
+      "compute_loss_in_ms": 285.45426247641444,
+      "epoch": 16.12,
+      "learning_rate/full": 0.00036401058098760525,
+      "loss": 3.0534,
+      "step": 39000,
+      "training_step_in_ms": 1088.1995187923312
+    },
+    {
+      "_prepare_inputs_in_ms": 4.224584739655256,
+      "compute_loss_in_ms": 285.29780930280685,
+      "epoch": 16.54,
+      "learning_rate/full": 0.00036219068645119566,
+      "loss": 3.0525,
+      "step": 40000,
+      "training_step_in_ms": 1088.9517585895956
+    },
+    {
+      "epoch": 16.54,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.0416412353515625,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.26328806809020683,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.03,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.271,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 40000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.210097938776016,
+      "compute_loss_in_ms": 285.2052926979959,
+      "epoch": 16.95,
+      "learning_rate/full": 0.0003603344533347134,
+      "loss": 3.0483,
+      "step": 41000,
+      "training_step_in_ms": 1085.6118382960558
+    },
+    {
+      "_prepare_inputs_in_ms": 4.248851552605629,
+      "compute_loss_in_ms": 285.6537539064884,
+      "epoch": 17.36,
+      "learning_rate/full": 0.00035843490089475537,
+      "loss": 3.0399,
+      "step": 42000,
+      "training_step_in_ms": 1091.6559825353324
+    },
+    {
+      "_prepare_inputs_in_ms": 4.214634284377098,
+      "compute_loss_in_ms": 285.3575124628842,
+      "epoch": 17.78,
+      "learning_rate/full": 0.0003564981368437495,
+      "loss": 3.0455,
+      "step": 43000,
+      "training_step_in_ms": 1088.8244492001832
+    },
+    {
+      "_prepare_inputs_in_ms": 4.242138650268316,
+      "compute_loss_in_ms": 285.5602181442082,
+      "epoch": 18.19,
+      "learning_rate/full": 0.00035452076268085417,
+      "loss": 3.0397,
+      "step": 44000,
+      "training_step_in_ms": 1087.6106830611825
+    },
+    {
+      "_prepare_inputs_in_ms": 4.231096193194389,
+      "compute_loss_in_ms": 285.38128500804305,
+      "epoch": 18.6,
+      "learning_rate/full": 0.0003525072209606466,
+      "loss": 3.0366,
+      "step": 45000,
+      "training_step_in_ms": 1090.0634618513286
+    },
+    {
+      "epoch": 18.6,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.0333669185638428,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.262616571295984,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.2623,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.389,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.231,
+      "step": 45000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.224474132061005,
+      "compute_loss_in_ms": 285.69229750707746,
+      "epoch": 19.02,
+      "learning_rate/full": 0.00035045605036568154,
+      "loss": 3.0403,
+      "step": 46000,
+      "training_step_in_ms": 1087.4972796961665
+    },
+    {
+      "_prepare_inputs_in_ms": 4.218343399465084,
+      "compute_loss_in_ms": 285.2950618080795,
+      "epoch": 19.43,
+      "learning_rate/full": 0.0003483677569916109,
+      "loss": 3.0296,
+      "step": 47000,
+      "training_step_in_ms": 1088.6641021184623
+    },
+    {
+      "_prepare_inputs_in_ms": 4.212014690041542,
+      "compute_loss_in_ms": 285.24483662098646,
+      "epoch": 19.84,
+      "learning_rate/full": 0.0003462450012513184,
+      "loss": 3.0351,
+      "step": 48000,
+      "training_step_in_ms": 1086.902916610241
+    },
+    {
+      "_prepare_inputs_in_ms": 4.228286672383547,
+      "compute_loss_in_ms": 285.4265847504139,
+      "epoch": 20.26,
+      "learning_rate/full": 0.0003440818719590809,
+      "loss": 3.0301,
+      "step": 49000,
+      "training_step_in_ms": 1089.4173335321248
+    },
+    {
+      "_prepare_inputs_in_ms": 4.209285493940115,
+      "compute_loss_in_ms": 285.17769135162234,
+      "epoch": 20.67,
+      "learning_rate/full": 0.0003418853377786221,
+      "loss": 3.0266,
+      "step": 50000,
+      "training_step_in_ms": 1092.610530115664
+    },
+    {
+      "epoch": 20.67,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.026047468185425,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.26601445767420673,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 111.3497,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.185,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.225,
+      "step": 50000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.215818756237263,
+      "compute_loss_in_ms": 285.3216738887131,
+      "epoch": 21.08,
+      "learning_rate/full": 0.00033965379551414244,
+      "loss": 3.0364,
+      "step": 51000,
+      "training_step_in_ms": 1090.6509163863957
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2158047296106815,
+      "compute_loss_in_ms": 285.3607781082392,
+      "epoch": 21.5,
+      "learning_rate/full": 0.00033738779576530426,
+      "loss": 3.0221,
+      "step": 52000,
+      "training_step_in_ms": 1089.9412010349333
+    },
+    {
+      "_prepare_inputs_in_ms": 4.225385930389166,
+      "compute_loss_in_ms": 285.4064598791301,
+      "epoch": 21.91,
+      "learning_rate/full": 0.0003350878976336386,
+      "loss": 3.0233,
+      "step": 53000,
+      "training_step_in_ms": 1086.4266870431602
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2348253689706326,
+      "compute_loss_in_ms": 285.5650148577988,
+      "epoch": 22.32,
+      "learning_rate/full": 0.0003327546685845955,
+      "loss": 3.0177,
+      "step": 54000,
+      "training_step_in_ms": 1090.5466065071523
+    },
+    {
+      "_prepare_inputs_in_ms": 4.244904510676861,
+      "compute_loss_in_ms": 285.4943734779954,
+      "epoch": 22.74,
+      "learning_rate/full": 0.00033038868430752995,
+      "loss": 3.0227,
+      "step": 55000,
+      "training_step_in_ms": 1089.5386388339102
+    },
+    {
+      "epoch": 22.74,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.0213677883148193,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.26657402454724916,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.863,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.282,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 55000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.222077999900027,
+      "compute_loss_in_ms": 285.3860865868628,
+      "epoch": 23.15,
+      "learning_rate/full": 0.00032798811209649607,
+      "loss": 3.0201,
+      "step": 56000,
+      "training_step_in_ms": 1089.8713997229934
+    },
+    {
+      "_prepare_inputs_in_ms": 4.218684710562229,
+      "compute_loss_in_ms": 285.196179587394,
+      "epoch": 23.56,
+      "learning_rate/full": 0.0003255583453025672,
+      "loss": 3.0133,
+      "step": 57000,
+      "training_step_in_ms": 1088.3847643770278
+    },
+    {
+      "_prepare_inputs_in_ms": 4.225137319415808,
+      "compute_loss_in_ms": 285.31357542052865,
+      "epoch": 23.98,
+      "learning_rate/full": 0.0003231000773635045,
+      "loss": 3.0174,
+      "step": 58000,
+      "training_step_in_ms": 1086.5370167195797
+    },
+    {
+      "_prepare_inputs_in_ms": 4.230918549001217,
+      "compute_loss_in_ms": 285.4190446138382,
+      "epoch": 24.39,
+      "learning_rate/full": 0.0003206140056326384,
+      "loss": 3.0116,
+      "step": 59000,
+      "training_step_in_ms": 1090.3938182927668
+    },
+    {
+      "_prepare_inputs_in_ms": 4.227691676467657,
+      "compute_loss_in_ms": 285.3117839321494,
+      "epoch": 24.8,
+      "learning_rate/full": 0.000318090679282307,
+      "loss": 3.0123,
+      "step": 60000,
+      "training_step_in_ms": 1090.3038867227733
+    },
+    {
+      "epoch": 24.8,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.0163190364837646,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2694944095513101,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.5302,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.304,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 60000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2251157578898635,
+      "compute_loss_in_ms": 285.44914393499494,
+      "epoch": 25.22,
+      "learning_rate/full": 0.0003155381574633497,
+      "loss": 3.0074,
+      "step": 61000,
+      "training_step_in_ms": 1087.7947441898286
+    },
+    {
+      "_prepare_inputs_in_ms": 4.233858399093151,
+      "compute_loss_in_ms": 285.5572083890438,
+      "epoch": 25.63,
+      "learning_rate/full": 0.0003129570712337902,
+      "loss": 3.0038,
+      "step": 62000,
+      "training_step_in_ms": 1091.406288355589
+    },
+    {
+      "_prepare_inputs_in_ms": 4.241201400756836,
+      "compute_loss_in_ms": 285.4022887274623,
+      "epoch": 26.04,
+      "learning_rate/full": 0.00031035068146119334,
+      "loss": 3.0069,
+      "step": 63000,
+      "training_step_in_ms": 1089.0374966450036
+    },
+    {
+      "_prepare_inputs_in_ms": 4.214517045766115,
+      "compute_loss_in_ms": 285.32751731202006,
+      "epoch": 26.46,
+      "learning_rate/full": 0.0003077170643091587,
+      "loss": 3.0004,
+      "step": 64000,
+      "training_step_in_ms": 1089.311513543129
+    },
+    {
+      "_prepare_inputs_in_ms": 4.223439604043961,
+      "compute_loss_in_ms": 285.6179902665317,
+      "epoch": 26.87,
+      "learning_rate/full": 0.00030505419362911944,
+      "loss": 3.0048,
+      "step": 65000,
+      "training_step_in_ms": 1088.4543421529233
+    },
+    {
+      "epoch": 26.87,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.012563705444336,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.26603180141607496,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.7667,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.355,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 65000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.235190596522354,
+      "compute_loss_in_ms": 285.59519398212433,
+      "epoch": 27.28,
+      "learning_rate/full": 0.0003023680520765336,
+      "loss": 2.9934,
+      "step": 66000,
+      "training_step_in_ms": 1088.0777766555548
+    },
+    {
+      "_prepare_inputs_in_ms": 4.217139046639204,
+      "compute_loss_in_ms": 285.50919711589813,
+      "epoch": 27.7,
+      "learning_rate/full": 0.0002996566527388639,
+      "loss": 2.9982,
+      "step": 67000,
+      "training_step_in_ms": 1091.6493426598608
+    },
+    {
+      "_prepare_inputs_in_ms": 4.244372218847275,
+      "compute_loss_in_ms": 285.5752951391041,
+      "epoch": 28.11,
+      "learning_rate/full": 0.0002969206646133254,
+      "loss": 2.9969,
+      "step": 68000,
+      "training_step_in_ms": 1088.2136982679367
+    },
+    {
+      "_prepare_inputs_in_ms": 4.213636931031942,
+      "compute_loss_in_ms": 285.255677562207,
+      "epoch": 28.52,
+      "learning_rate/full": 0.0002941607627640486,
+      "loss": 2.9923,
+      "step": 69000,
+      "training_step_in_ms": 1088.9643149748445
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2318920604884624,
+      "compute_loss_in_ms": 285.6222639977932,
+      "epoch": 28.94,
+      "learning_rate/full": 0.0002913748308243434,
+      "loss": 2.9912,
+      "step": 70000,
+      "training_step_in_ms": 1086.640508864075
+    },
+    {
+      "epoch": 28.94,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.01218318939209,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2704365100152127,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.8069,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.286,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 70000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.22964243627176,
+      "compute_loss_in_ms": 285.81058219075203,
+      "epoch": 29.35,
+      "learning_rate/full": 0.00028856630835486283,
+      "loss": 2.9907,
+      "step": 71000,
+      "training_step_in_ms": 1089.0001546032727
+    },
+    {
+      "_prepare_inputs_in_ms": 4.219989389181137,
+      "compute_loss_in_ms": 285.4850408025086,
+      "epoch": 29.76,
+      "learning_rate/full": 0.00028574157192993993,
+      "loss": 2.9922,
+      "step": 72000,
+      "training_step_in_ms": 1112.2257943935692
+    },
+    {
+      "_prepare_inputs_in_ms": 4.236269619315863,
+      "compute_loss_in_ms": 285.3517268039286,
+      "epoch": 30.18,
+      "learning_rate/full": 0.0002828899985518552,
+      "loss": 2.9829,
+      "step": 73000,
+      "training_step_in_ms": 1115.709298092872
+    },
+    {
+      "_prepare_inputs_in_ms": 4.212801028043032,
+      "compute_loss_in_ms": 285.1435379870236,
+      "epoch": 30.59,
+      "learning_rate/full": 0.0002800179323426103,
+      "loss": 2.9854,
+      "step": 74000,
+      "training_step_in_ms": 1112.750349264592
+    },
+    {
+      "_prepare_inputs_in_ms": 4.226626381278038,
+      "compute_loss_in_ms": 285.37043143063784,
+      "epoch": 31.0,
+      "learning_rate/full": 0.0002771289848538608,
+      "loss": 2.9928,
+      "step": 75000,
+      "training_step_in_ms": 1113.63447811082
+    },
+    {
+      "epoch": 31.0,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.00150203704834,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27227553064507803,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.2006,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.259,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 75000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.221666444365571,
+      "compute_loss_in_ms": 285.4123991020024,
+      "epoch": 31.42,
+      "learning_rate/full": 0.00027422392710754273,
+      "loss": 2.9787,
+      "step": 76000,
+      "training_step_in_ms": 1110.9835148528218
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2181270979344845,
+      "compute_loss_in_ms": 285.3203030079603,
+      "epoch": 31.83,
+      "learning_rate/full": 0.0002712947161076778,
+      "loss": 2.9822,
+      "step": 77000,
+      "training_step_in_ms": 1113.9973731786013
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2350912764668465,
+      "compute_loss_in_ms": 285.4480539858341,
+      "epoch": 32.24,
+      "learning_rate/full": 0.00026835083436875734,
+      "loss": 2.9765,
+      "step": 78000,
+      "training_step_in_ms": 1115.4474330842495
+    },
+    {
+      "_prepare_inputs_in_ms": 4.229009635746479,
+      "compute_loss_in_ms": 285.5191092900932,
+      "epoch": 32.66,
+      "learning_rate/full": 0.0002653871161688328,
+      "loss": 2.9801,
+      "step": 79000,
+      "training_step_in_ms": 1113.7964499779046
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2362766563892365,
+      "compute_loss_in_ms": 285.7597692273557,
+      "epoch": 33.07,
+      "learning_rate/full": 0.00026241022007566643,
+      "loss": 2.9807,
+      "step": 80000,
+      "training_step_in_ms": 1115.4754909984767
+    },
+    {
+      "epoch": 33.07,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 3.000786781311035,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2716794971860456,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.7776,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.354,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 80000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.232099074416045,
+      "compute_loss_in_ms": 285.50172889232635,
+      "epoch": 33.48,
+      "learning_rate/full": 0.0002594179251945605,
+      "loss": 2.9739,
+      "step": 81000,
+      "training_step_in_ms": 1110.9056022837758
+    },
+    {
+      "_prepare_inputs_in_ms": 4.237273696810007,
+      "compute_loss_in_ms": 285.7116014882922,
+      "epoch": 33.9,
+      "learning_rate/full": 0.00025641096982950234,
+      "loss": 2.9746,
+      "step": 82000,
+      "training_step_in_ms": 1110.9737426675856
+    },
+    {
+      "_prepare_inputs_in_ms": 4.237969063222408,
+      "compute_loss_in_ms": 285.6498990356922,
+      "epoch": 34.31,
+      "learning_rate/full": 0.00025339009590173424,
+      "loss": 2.9727,
+      "step": 83000,
+      "training_step_in_ms": 1117.3185790739954
+    },
+    {
+      "_prepare_inputs_in_ms": 4.234774090349674,
+      "compute_loss_in_ms": 285.5246250964701,
+      "epoch": 34.73,
+      "learning_rate/full": 0.00025035604876669546,
+      "loss": 2.9709,
+      "step": 84000,
+      "training_step_in_ms": 1111.8130441047251
+    },
+    {
+      "_prepare_inputs_in_ms": 4.248618151992559,
+      "compute_loss_in_ms": 285.6561874523759,
+      "epoch": 35.14,
+      "learning_rate/full": 0.00024731263251348453,
+      "loss": 2.969,
+      "step": 85000,
+      "training_step_in_ms": 1112.8755748830736
+    },
+    {
+      "epoch": 35.14,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9980032444000244,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2741870945987276,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.6613,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.229,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.226,
+      "step": 85000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.240446417796902,
+      "compute_loss_in_ms": 285.5394543148577,
+      "epoch": 35.55,
+      "learning_rate/full": 0.00024425143236331536,
+      "loss": 2.9639,
+      "step": 86000,
+      "training_step_in_ms": 1110.7865899279714
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2231163419783115,
+      "compute_loss_in_ms": 285.532758615911,
+      "epoch": 35.97,
+      "learning_rate/full": 0.000241179291965253,
+      "loss": 2.971,
+      "step": 87000,
+      "training_step_in_ms": 1113.825252827257
+    },
+    {
+      "_prepare_inputs_in_ms": 4.229363452643156,
+      "compute_loss_in_ms": 285.72211230918765,
+      "epoch": 36.38,
+      "learning_rate/full": 0.0002381000579951894,
+      "loss": 2.9636,
+      "step": 88000,
+      "training_step_in_ms": 1118.3622099086642
+    },
+    {
+      "_prepare_inputs_in_ms": 4.2257860116660595,
+      "compute_loss_in_ms": 285.5116978622973,
+      "epoch": 36.79,
+      "learning_rate/full": 0.00023501142340591894,
+      "loss": 2.9656,
+      "step": 89000,
+      "training_step_in_ms": 1111.638593826443
+    },
+    {
+      "_prepare_inputs_in_ms": 4.221606273204088,
+      "compute_loss_in_ms": 285.61226362735033,
+      "epoch": 37.21,
+      "learning_rate/full": 0.00023191415027181022,
+      "loss": 2.9615,
+      "step": 90000,
+      "training_step_in_ms": 1116.2303377054632
+    },
+    {
+      "epoch": 37.21,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.996258020401001,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27303322222107285,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.8741,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.348,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 90000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.207164986831386,
+      "compute_loss_in_ms": 285.1755935549736,
+      "epoch": 37.62,
+      "learning_rate/full": 0.00022881211473645583,
+      "loss": 2.963,
+      "step": 91000,
+      "training_step_in_ms": 1110.297369044274
+    },
+    {
+      "_prepare_inputs_in_ms": 4.222084645181894,
+      "compute_loss_in_ms": 285.53527039662004,
+      "epoch": 38.03,
+      "learning_rate/full": 0.00022570298446764845,
+      "loss": 2.9633,
+      "step": 92000,
+      "training_step_in_ms": 1112.9915070161223
+    },
+    {
+      "_prepare_inputs_in_ms": 4.232885275036097,
+      "compute_loss_in_ms": 285.5158912166953,
+      "epoch": 38.45,
+      "learning_rate/full": 0.00022258127581536945,
+      "loss": 2.9538,
+      "step": 93000,
+      "training_step_in_ms": 1114.245859079063
+    },
+    {
+      "_prepare_inputs_in_ms": 4.250343676656485,
+      "compute_loss_in_ms": 285.64605471119285,
+      "epoch": 38.86,
+      "learning_rate/full": 0.00021945398441148287,
+      "loss": 2.9572,
+      "step": 94000,
+      "training_step_in_ms": 1112.7353053241968
+    },
+    {
+      "_prepare_inputs_in_ms": 4.230448927730322,
+      "compute_loss_in_ms": 285.5413333699107,
+      "epoch": 39.27,
+      "learning_rate/full": 0.00021632501765960936,
+      "loss": 2.958,
+      "step": 95000,
+      "training_step_in_ms": 1112.5333589836955
+    },
+    {
+      "epoch": 39.27,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.993772506713867,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27348855682421375,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.9695,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.342,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.229,
+      "step": 95000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.215012745159428,
+      "compute_loss_in_ms": 285.35930648073554,
+      "epoch": 39.69,
+      "learning_rate/full": 0.0002131920229539048,
+      "loss": 2.9558,
+      "step": 96000,
+      "training_step_in_ms": 1111.498819194734
+    },
+    {
+      "_prepare_inputs_in_ms": 4.221161104738712,
+      "compute_loss_in_ms": 285.453462138772,
+      "epoch": 40.1,
+      "learning_rate/full": 0.00021005263255270636,
+      "loss": 2.9559,
+      "step": 97000,
+      "training_step_in_ms": 1114.1434171833098
+    },
+    {
+      "_prepare_inputs_in_ms": 4.220512144267559,
+      "compute_loss_in_ms": 285.53688745573163,
+      "epoch": 40.51,
+      "learning_rate/full": 0.0002069107568468244,
+      "loss": 2.9525,
+      "step": 98000,
+      "training_step_in_ms": 1116.582923579961
+    },
+    {
+      "_prepare_inputs_in_ms": 4.22445010766387,
+      "compute_loss_in_ms": 285.29584189876914,
+      "epoch": 40.93,
+      "learning_rate/full": 0.00020377031677881017,
+      "loss": 2.9509,
+      "step": 99000,
+      "training_step_in_ms": 1112.9648886173964
+    },
+    {
+      "_prepare_inputs_in_ms": 4.237125843763351,
+      "compute_loss_in_ms": 285.87847367301583,
+      "epoch": 41.34,
+      "learning_rate/full": 0.00020062580171962844,
+      "loss": 2.9397,
+      "step": 100000,
+      "training_step_in_ms": 1119.0427548959851
+    },
+    {
+      "epoch": 41.34,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9917938709259033,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27402243679815436,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.6022,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.366,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 100000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.220401306704777,
+      "compute_loss_in_ms": 285.4624089188874,
+      "epoch": 41.75,
+      "learning_rate/full": 0.00019748427643225179,
+      "loss": 2.9497,
+      "step": 101000,
+      "training_step_in_ms": 1113.234252423048
+    },
+    {
+      "_prepare_inputs_in_ms": 4.218052037060261,
+      "compute_loss_in_ms": 285.50472677126527,
+      "epoch": 42.17,
+      "learning_rate/full": 0.0001943402283833764,
+      "loss": 2.9442,
+      "step": 102000,
+      "training_step_in_ms": 1114.7688954658806
+    },
+    {
+      "_prepare_inputs_in_ms": 4.22609718888998,
+      "compute_loss_in_ms": 285.6670557744801,
+      "epoch": 42.58,
+      "learning_rate/full": 0.0001911975795955237,
+      "loss": 2.938,
+      "step": 103000,
+      "training_step_in_ms": 1114.2966277077794
+    },
+    {
+      "_prepare_inputs_in_ms": 4.231020983308554,
+      "compute_loss_in_ms": 285.41950649395585,
+      "epoch": 42.99,
+      "learning_rate/full": 0.00018806024615043859,
+      "loss": 2.9429,
+      "step": 104000,
+      "training_step_in_ms": 1112.3123243488371
+    },
+    {
+      "_prepare_inputs_in_ms": 4.231835335493088,
+      "compute_loss_in_ms": 285.55181711539626,
+      "epoch": 43.41,
+      "learning_rate/full": 0.0001849258586609575,
+      "loss": 2.9355,
+      "step": 105000,
+      "training_step_in_ms": 1116.4177654609084
+    },
+    {
+      "epoch": 43.41,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9879465103149414,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.275499690988243,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.1682,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.328,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.229,
+      "step": 105000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.192932469815743,
+      "compute_loss_in_ms": 285.3939059227705,
+      "epoch": 43.82,
+      "learning_rate/full": 0.0001817951904902926,
+      "loss": 2.9411,
+      "step": 106000,
+      "training_step_in_ms": 1111.0442412495613
+    },
+    {
+      "_prepare_inputs_in_ms": 4.18786546587944,
+      "compute_loss_in_ms": 285.60110822692513,
+      "epoch": 44.23,
+      "learning_rate/full": 0.00017866588728649688,
+      "loss": 2.9389,
+      "step": 107000,
+      "training_step_in_ms": 1114.2360820770264
+    },
+    {
+      "_prepare_inputs_in_ms": 4.191282417625189,
+      "compute_loss_in_ms": 285.483426745981,
+      "epoch": 44.65,
+      "learning_rate/full": 0.00017554497963615946,
+      "loss": 2.9396,
+      "step": 108000,
+      "training_step_in_ms": 1113.102070134133
+    },
+    {
+      "_prepare_inputs_in_ms": 4.215992953628302,
+      "compute_loss_in_ms": 285.87069864198565,
+      "epoch": 45.06,
+      "learning_rate/full": 0.00017243010589616854,
+      "loss": 2.9398,
+      "step": 109000,
+      "training_step_in_ms": 1115.9531138837337
+    },
+    {
+      "_prepare_inputs_in_ms": 4.180787291377783,
+      "compute_loss_in_ms": 285.6199772916734,
+      "epoch": 45.47,
+      "learning_rate/full": 0.00016932203461501055,
+      "loss": 2.9354,
+      "step": 110000,
+      "training_step_in_ms": 1112.3631411641836
+    },
+    {
+      "epoch": 45.47,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9866650104522705,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27461185232552654,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.9151,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.345,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 110000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1715224159926905,
+      "compute_loss_in_ms": 285.4828831627965,
+      "epoch": 45.89,
+      "learning_rate/full": 0.00016622153266276704,
+      "loss": 2.9325,
+      "step": 111000,
+      "training_step_in_ms": 1110.1254166848958
+    },
+    {
+      "_prepare_inputs_in_ms": 4.183446723967791,
+      "compute_loss_in_ms": 285.5392692387104,
+      "epoch": 46.3,
+      "learning_rate/full": 0.00016312936504190095,
+      "loss": 2.9271,
+      "step": 112000,
+      "training_step_in_ms": 1116.0438286960125
+    },
+    {
+      "_prepare_inputs_in_ms": 4.181461203843355,
+      "compute_loss_in_ms": 285.49605195596814,
+      "epoch": 46.71,
+      "learning_rate/full": 0.00016004321335415234,
+      "loss": 2.9338,
+      "step": 113000,
+      "training_step_in_ms": 1114.0362310223281
+    },
+    {
+      "_prepare_inputs_in_ms": 4.20052033662796,
+      "compute_loss_in_ms": 285.6174496598542,
+      "epoch": 47.13,
+      "learning_rate/full": 0.00015696694015230966,
+      "loss": 2.93,
+      "step": 114000,
+      "training_step_in_ms": 1112.158472020179
+    },
+    {
+      "_prepare_inputs_in_ms": 4.184404268860817,
+      "compute_loss_in_ms": 285.52299703657627,
+      "epoch": 47.54,
+      "learning_rate/full": 0.00015390436604853944,
+      "loss": 2.9265,
+      "step": 115000,
+      "training_step_in_ms": 1114.8844772167504
+    },
+    {
+      "epoch": 47.54,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.985903263092041,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2781795787221218,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.8815,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.281,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 115000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.152440788542352,
+      "compute_loss_in_ms": 285.35162526741624,
+      "epoch": 47.95,
+      "learning_rate/full": 0.00015085316535270307,
+      "loss": 2.9288,
+      "step": 116000,
+      "training_step_in_ms": 1108.712267011404
+    },
+    {
+      "_prepare_inputs_in_ms": 4.172631837427616,
+      "compute_loss_in_ms": 285.6740382388234,
+      "epoch": 48.37,
+      "learning_rate/full": 0.0001478110551124508,
+      "loss": 2.9217,
+      "step": 117000,
+      "training_step_in_ms": 1119.9308814108372
+    },
+    {
+      "_prepare_inputs_in_ms": 4.173140484839678,
+      "compute_loss_in_ms": 285.33917328342795,
+      "epoch": 48.78,
+      "learning_rate/full": 0.0001447878925453241,
+      "loss": 2.9189,
+      "step": 118000,
+      "training_step_in_ms": 1114.4561827853322
+    },
+    {
+      "_prepare_inputs_in_ms": 4.208178836852312,
+      "compute_loss_in_ms": 285.7572955302894,
+      "epoch": 49.19,
+      "learning_rate/full": 0.00014177531694909012,
+      "loss": 2.9234,
+      "step": 119000,
+      "training_step_in_ms": 1117.5116944983602
+    },
+    {
+      "_prepare_inputs_in_ms": 4.171714887022972,
+      "compute_loss_in_ms": 285.51325725764036,
+      "epoch": 49.61,
+      "learning_rate/full": 0.00013877411365635932,
+      "loss": 2.919,
+      "step": 120000,
+      "training_step_in_ms": 1114.4923375099897
+    },
+    {
+      "epoch": 49.61,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.984545946121216,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27571246372317004,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 107.5626,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.438,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.232,
+      "step": 120000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.169439662520478,
+      "compute_loss_in_ms": 285.2081711217761,
+      "epoch": 50.02,
+      "learning_rate/full": 0.00013579102545242515,
+      "loss": 2.9214,
+      "step": 121000,
+      "training_step_in_ms": 1117.0064120963216
+    },
+    {
+      "_prepare_inputs_in_ms": 4.175419889390469,
+      "compute_loss_in_ms": 285.44106700643897,
+      "epoch": 50.43,
+      "learning_rate/full": 0.00013282377985218108,
+      "loss": 2.9123,
+      "step": 122000,
+      "training_step_in_ms": 1117.1209677942097
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1646773256361485,
+      "compute_loss_in_ms": 285.5908838920295,
+      "epoch": 50.85,
+      "learning_rate/full": 0.0001298731089790791,
+      "loss": 2.9181,
+      "step": 123000,
+      "training_step_in_ms": 1114.7672307156026
+    },
+    {
+      "_prepare_inputs_in_ms": 4.172112949192524,
+      "compute_loss_in_ms": 285.507780585438,
+      "epoch": 51.26,
+      "learning_rate/full": 0.0001269397408670054,
+      "loss": 2.9124,
+      "step": 124000,
+      "training_step_in_ms": 1116.1798403412104
+    },
+    {
+      "_prepare_inputs_in_ms": 4.171217355877161,
+      "compute_loss_in_ms": 285.5308585166931,
+      "epoch": 51.67,
+      "learning_rate/full": 0.0001240214902928718,
+      "loss": 2.9169,
+      "step": 125000,
+      "training_step_in_ms": 1119.0364625044167
+    },
+    {
+      "epoch": 51.67,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9826271533966064,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27485379257600506,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.595,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.3,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 125000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1942793857760545,
+      "compute_loss_in_ms": 285.4048260115087,
+      "epoch": 52.09,
+      "learning_rate/full": 0.00012112491367260039,
+      "loss": 2.9117,
+      "step": 126000,
+      "training_step_in_ms": 1113.5866565182805
+    },
+    {
+      "_prepare_inputs_in_ms": 4.165726162493229,
+      "compute_loss_in_ms": 285.52386473864317,
+      "epoch": 52.5,
+      "learning_rate/full": 0.0001182477982996471,
+      "loss": 2.9123,
+      "step": 127000,
+      "training_step_in_ms": 1115.4936682023108
+    },
+    {
+      "_prepare_inputs_in_ms": 4.165291707962751,
+      "compute_loss_in_ms": 285.37697672098875,
+      "epoch": 52.91,
+      "learning_rate/full": 0.00011539085405917883,
+      "loss": 2.9127,
+      "step": 128000,
+      "training_step_in_ms": 1112.6525225900114
+    },
+    {
+      "_prepare_inputs_in_ms": 4.205962881445885,
+      "compute_loss_in_ms": 285.6344051398337,
+      "epoch": 53.33,
+      "learning_rate/full": 0.00011255195764553374,
+      "loss": 2.9072,
+      "step": 129000,
+      "training_step_in_ms": 1114.8041105866432
+    },
+    {
+      "_prepare_inputs_in_ms": 4.192876800894737,
+      "compute_loss_in_ms": 285.6524411961436,
+      "epoch": 53.74,
+      "learning_rate/full": 0.00010973748719012139,
+      "loss": 2.9105,
+      "step": 130000,
+      "training_step_in_ms": 1113.7267719507217
+    },
+    {
+      "epoch": 53.74,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.983597993850708,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2789962669075427,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.5326,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.304,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 130000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.207222261806813,
+      "compute_loss_in_ms": 285.5676299482584,
+      "epoch": 54.15,
+      "learning_rate/full": 0.0001069452876591179,
+      "loss": 2.9063,
+      "step": 131000,
+      "training_step_in_ms": 1114.3601090423763
+    },
+    {
+      "_prepare_inputs_in_ms": 4.189547453075647,
+      "compute_loss_in_ms": 285.45121479034424,
+      "epoch": 54.57,
+      "learning_rate/full": 0.00010417604798597693,
+      "loss": 2.9068,
+      "step": 132000,
+      "training_step_in_ms": 1116.6341638937593
+    },
+    {
+      "_prepare_inputs_in_ms": 4.203598700463772,
+      "compute_loss_in_ms": 285.4171659834683,
+      "epoch": 54.98,
+      "learning_rate/full": 0.00010142771516637335,
+      "loss": 2.8989,
+      "step": 133000,
+      "training_step_in_ms": 1110.624721519649
+    },
+    {
+      "_prepare_inputs_in_ms": 4.217664018273354,
+      "compute_loss_in_ms": 285.5546323284507,
+      "epoch": 55.39,
+      "learning_rate/full": 9.870646386303746e-05,
+      "loss": 2.8974,
+      "step": 134000,
+      "training_step_in_ms": 1115.4996632412076
+    },
+    {
+      "_prepare_inputs_in_ms": 4.18796281516552,
+      "compute_loss_in_ms": 285.3805873543024,
+      "epoch": 55.81,
+      "learning_rate/full": 9.601020522405566e-05,
+      "loss": 2.8997,
+      "step": 135000,
+      "training_step_in_ms": 1114.3139710351825
+    },
+    {
+      "epoch": 55.81,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9829113483428955,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2781282773190371,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.4066,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.312,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.229,
+      "step": 135000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.179593571802465,
+      "compute_loss_in_ms": 285.6101936176419,
+      "epoch": 56.22,
+      "learning_rate/full": 9.333960451090202e-05,
+      "loss": 2.9042,
+      "step": 136000,
+      "training_step_in_ms": 1112.9348441772163
+    },
+    {
+      "_prepare_inputs_in_ms": 4.178862765431404,
+      "compute_loss_in_ms": 285.4050006046891,
+      "epoch": 56.63,
+      "learning_rate/full": 9.069532065434167e-05,
+      "loss": 2.8997,
+      "step": 137000,
+      "training_step_in_ms": 1110.1743725985289
+    },
+    {
+      "_prepare_inputs_in_ms": 4.17331463098526,
+      "compute_loss_in_ms": 285.60341618955135,
+      "epoch": 57.05,
+      "learning_rate/full": 8.807539988537217e-05,
+      "loss": 2.8999,
+      "step": 138000,
+      "training_step_in_ms": 1112.455148395151
+    },
+    {
+      "_prepare_inputs_in_ms": 4.164013888686895,
+      "compute_loss_in_ms": 285.4794158451259,
+      "epoch": 57.46,
+      "learning_rate/full": 8.54857283641461e-05,
+      "loss": 2.9008,
+      "step": 139000,
+      "training_step_in_ms": 1112.6597697511315
+    },
+    {
+      "_prepare_inputs_in_ms": 4.174750838428736,
+      "compute_loss_in_ms": 285.3338685967028,
+      "epoch": 57.88,
+      "learning_rate/full": 8.292176191253292e-05,
+      "loss": 2.9003,
+      "step": 140000,
+      "training_step_in_ms": 1112.5756445713341
+    },
+    {
+      "epoch": 57.88,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.983377456665039,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2765083682413393,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.2039,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.326,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.229,
+      "step": 140000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.186901926267438,
+      "compute_loss_in_ms": 285.5366085804999,
+      "epoch": 58.29,
+      "learning_rate/full": 8.038926097845864e-05,
+      "loss": 2.89,
+      "step": 141000,
+      "training_step_in_ms": 1111.180388186127
+    },
+    {
+      "_prepare_inputs_in_ms": 4.166560925543308,
+      "compute_loss_in_ms": 285.2483623623848,
+      "epoch": 58.7,
+      "learning_rate/full": 7.788627220446403e-05,
+      "loss": 2.8908,
+      "step": 142000,
+      "training_step_in_ms": 1112.5438826270401
+    },
+    {
+      "_prepare_inputs_in_ms": 4.172076798975468,
+      "compute_loss_in_ms": 285.31228306889534,
+      "epoch": 59.12,
+      "learning_rate/full": 7.541341316557602e-05,
+      "loss": 2.888,
+      "step": 143000,
+      "training_step_in_ms": 1114.1774371489882
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1760097071528435,
+      "compute_loss_in_ms": 285.51259553432465,
+      "epoch": 59.53,
+      "learning_rate/full": 7.29712940027603e-05,
+      "loss": 2.8933,
+      "step": 144000,
+      "training_step_in_ms": 1113.300205629319
+    },
+    {
+      "_prepare_inputs_in_ms": 4.171927604824305,
+      "compute_loss_in_ms": 285.47950995340943,
+      "epoch": 59.94,
+      "learning_rate/full": 7.0560517272378e-05,
+      "loss": 2.8946,
+      "step": 145000,
+      "training_step_in_ms": 1110.925380833447
+    },
+    {
+      "epoch": 59.94,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9790754318237305,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.27903548274273476,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 108.6903,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.36,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.23,
+      "step": 145000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.179818335829712,
+      "compute_loss_in_ms": 285.67399540916085,
+      "epoch": 60.36,
+      "learning_rate/full": 6.818167779751427e-05,
+      "loss": 2.8868,
+      "step": 146000,
+      "training_step_in_ms": 1113.526238951832
+    },
+    {
+      "_prepare_inputs_in_ms": 4.174669615924358,
+      "compute_loss_in_ms": 285.47526767477393,
+      "epoch": 60.77,
+      "learning_rate/full": 6.58353625212141e-05,
+      "loss": 2.8908,
+      "step": 147000,
+      "training_step_in_ms": 1112.4517313353717
+    },
+    {
+      "_prepare_inputs_in_ms": 4.206293076276779,
+      "compute_loss_in_ms": 285.6568570397794,
+      "epoch": 61.18,
+      "learning_rate/full": 6.351985161010259e-05,
+      "loss": 2.886,
+      "step": 148000,
+      "training_step_in_ms": 1114.374936837703
+    },
+    {
+      "_prepare_inputs_in_ms": 4.173767436295748,
+      "compute_loss_in_ms": 285.36785116791725,
+      "epoch": 61.6,
+      "learning_rate/full": 6.124034730854495e-05,
+      "loss": 2.8877,
+      "step": 149000,
+      "training_step_in_ms": 1113.2941167131066
+    },
+    {
+      "_prepare_inputs_in_ms": 4.186146479099989,
+      "compute_loss_in_ms": 285.7034795098007,
+      "epoch": 62.01,
+      "learning_rate/full": 5.8995079874983696e-05,
+      "loss": 2.887,
+      "step": 150000,
+      "training_step_in_ms": 1114.685032505542
+    },
+    {
+      "epoch": 62.01,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9789364337921143,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2814554306931564,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.1484,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.263,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 150000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.176326023369301,
+      "compute_loss_in_ms": 285.46767780929804,
+      "epoch": 62.42,
+      "learning_rate/full": 5.678240821954202e-05,
+      "loss": 2.8806,
+      "step": 151000,
+      "training_step_in_ms": 1111.9945207312703
+    },
+    {
+      "_prepare_inputs_in_ms": 4.179010454565287,
+      "compute_loss_in_ms": 285.4651117064059,
+      "epoch": 62.84,
+      "learning_rate/full": 5.460730353825116e-05,
+      "loss": 2.8834,
+      "step": 152000,
+      "training_step_in_ms": 1110.2052029296756
+    },
+    {
+      "_prepare_inputs_in_ms": 4.173945639282465,
+      "compute_loss_in_ms": 285.60662161558867,
+      "epoch": 63.25,
+      "learning_rate/full": 5.246594910646354e-05,
+      "loss": 2.8816,
+      "step": 153000,
+      "training_step_in_ms": 1115.5096570029855
+    },
+    {
+      "_prepare_inputs_in_ms": 4.187856215983629,
+      "compute_loss_in_ms": 285.6514365822077,
+      "epoch": 63.66,
+      "learning_rate/full": 5.0363155892235236e-05,
+      "loss": 2.8866,
+      "step": 154000,
+      "training_step_in_ms": 1113.550076983869
+    },
+    {
+      "_prepare_inputs_in_ms": 4.191848460584879,
+      "compute_loss_in_ms": 285.825974162668,
+      "epoch": 64.08,
+      "learning_rate/full": 4.8297283330226226e-05,
+      "loss": 2.8812,
+      "step": 155000,
+      "training_step_in_ms": 1113.2165458351374
+    },
+    {
+      "epoch": 64.08,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9803640842437744,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.28091860195000073,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.2234,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.258,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 155000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.170489263970677,
+      "compute_loss_in_ms": 285.33177894353867,
+      "epoch": 64.49,
+      "learning_rate/full": 4.626682959161812e-05,
+      "loss": 2.8769,
+      "step": 156000,
+      "training_step_in_ms": 1107.8080112151802
+    },
+    {
+      "_prepare_inputs_in_ms": 4.165396176278591,
+      "compute_loss_in_ms": 285.30143217742443,
+      "epoch": 64.9,
+      "learning_rate/full": 4.427635648454991e-05,
+      "loss": 2.8775,
+      "step": 157000,
+      "training_step_in_ms": 1111.1020593941212
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1787227392196655,
+      "compute_loss_in_ms": 285.6033191792667,
+      "epoch": 65.32,
+      "learning_rate/full": 4.2324305856376166e-05,
+      "loss": 2.8824,
+      "step": 158000,
+      "training_step_in_ms": 1114.8409751541913
+    },
+    {
+      "_prepare_inputs_in_ms": 4.168594349175692,
+      "compute_loss_in_ms": 285.385168325156,
+      "epoch": 65.73,
+      "learning_rate/full": 4.040926393437829e-05,
+      "loss": 2.8742,
+      "step": 159000,
+      "training_step_in_ms": 1112.1875176765025
+    },
+    {
+      "_prepare_inputs_in_ms": 4.188820585608482,
+      "compute_loss_in_ms": 285.74348379299045,
+      "epoch": 66.14,
+      "learning_rate/full": 3.853553323166454e-05,
+      "loss": 2.874,
+      "step": 160000,
+      "training_step_in_ms": 1114.071101732552
+    },
+    {
+      "epoch": 66.14,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.980635404586792,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2805308037735457,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.2483,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.323,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.229,
+      "step": 160000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.17639887550982,
+      "compute_loss_in_ms": 285.54739573970437,
+      "epoch": 66.56,
+      "learning_rate/full": 3.6703457126541777e-05,
+      "loss": 2.8751,
+      "step": 161000,
+      "training_step_in_ms": 1101.9119679294527
+    },
+    {
+      "_prepare_inputs_in_ms": 4.183506786823273,
+      "compute_loss_in_ms": 285.5607514716685,
+      "epoch": 66.97,
+      "learning_rate/full": 3.4908041134979454e-05,
+      "loss": 2.8781,
+      "step": 162000,
+      "training_step_in_ms": 1103.5643678978086
+    },
+    {
+      "_prepare_inputs_in_ms": 4.188889868557453,
+      "compute_loss_in_ms": 285.67521207407117,
+      "epoch": 67.38,
+      "learning_rate/full": 3.315517477040358e-05,
+      "loss": 2.878,
+      "step": 163000,
+      "training_step_in_ms": 1105.4489219635725
+    },
+    {
+      "_prepare_inputs_in_ms": 4.175573732703924,
+      "compute_loss_in_ms": 285.59245705604553,
+      "epoch": 67.8,
+      "learning_rate/full": 3.14417822230312e-05,
+      "loss": 2.8757,
+      "step": 164000,
+      "training_step_in_ms": 1105.2035297378898
+    },
+    {
+      "_prepare_inputs_in_ms": 4.187327720224857,
+      "compute_loss_in_ms": 285.88794915005565,
+      "epoch": 68.21,
+      "learning_rate/full": 2.9771712961539955e-05,
+      "loss": 2.8763,
+      "step": 165000,
+      "training_step_in_ms": 1105.2709091752768
+    },
+    {
+      "epoch": 68.21,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9800891876220703,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2802985242098701,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.8323,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.218,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.226,
+      "step": 165000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.207263810605538,
+      "compute_loss_in_ms": 285.7474652826786,
+      "epoch": 68.62,
+      "learning_rate/full": 2.8143644982694906e-05,
+      "loss": 2.8761,
+      "step": 166000,
+      "training_step_in_ms": 1100.3973172418773
+    },
+    {
+      "_prepare_inputs_in_ms": 4.210827711969614,
+      "compute_loss_in_ms": 285.5911776944995,
+      "epoch": 69.04,
+      "learning_rate/full": 2.65579799879085e-05,
+      "loss": 2.8751,
+      "step": 167000,
+      "training_step_in_ms": 1104.1183153651655
+    },
+    {
+      "_prepare_inputs_in_ms": 4.198013573884964,
+      "compute_loss_in_ms": 285.68624898046255,
+      "epoch": 69.45,
+      "learning_rate/full": 2.5015109216291467e-05,
+      "loss": 2.8722,
+      "step": 168000,
+      "training_step_in_ms": 1105.6054084450006
+    },
+    {
+      "_prepare_inputs_in_ms": 4.18689937889576,
+      "compute_loss_in_ms": 285.61638662964106,
+      "epoch": 69.86,
+      "learning_rate/full": 2.3515413348120198e-05,
+      "loss": 2.8743,
+      "step": 169000,
+      "training_step_in_ms": 1102.9215082861483
+    },
+    {
+      "_prepare_inputs_in_ms": 4.210809834301472,
+      "compute_loss_in_ms": 285.83876856043935,
+      "epoch": 70.28,
+      "learning_rate/full": 2.20592624109097e-05,
+      "loss": 2.8696,
+      "step": 170000,
+      "training_step_in_ms": 1106.3744595497847
+    },
+    {
+      "epoch": 70.28,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.979834794998169,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.28091535007381413,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.6944,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.293,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 170000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.173562210507509,
+      "compute_loss_in_ms": 285.2328538559377,
+      "epoch": 70.69,
+      "learning_rate/full": 2.064562414115867e-05,
+      "loss": 2.8723,
+      "step": 171000,
+      "training_step_in_ms": 1101.5736067220569
+    },
+    {
+      "_prepare_inputs_in_ms": 4.184829708188772,
+      "compute_loss_in_ms": 285.7893420346081,
+      "epoch": 71.1,
+      "learning_rate/full": 1.9277674551421355e-05,
+      "loss": 2.8712,
+      "step": 172000,
+      "training_step_in_ms": 1104.9391337744892
+    },
+    {
+      "_prepare_inputs_in_ms": 4.175479732453823,
+      "compute_loss_in_ms": 285.58204352483153,
+      "epoch": 71.52,
+      "learning_rate/full": 1.7954315491282236e-05,
+      "loss": 2.868,
+      "step": 173000,
+      "training_step_in_ms": 1106.567913543433
+    },
+    {
+      "_prepare_inputs_in_ms": 4.216715902090073,
+      "compute_loss_in_ms": 285.6606830134988,
+      "epoch": 71.93,
+      "learning_rate/full": 1.6674616367900976e-05,
+      "loss": 2.8675,
+      "step": 174000,
+      "training_step_in_ms": 1104.7934159226716
+    },
+    {
+      "_prepare_inputs_in_ms": 4.171092137694359,
+      "compute_loss_in_ms": 285.4122787192464,
+      "epoch": 72.34,
+      "learning_rate/full": 1.5441452273561308e-05,
+      "loss": 2.8671,
+      "step": 175000,
+      "training_step_in_ms": 1106.7299974374473
+    },
+    {
+      "epoch": 72.34,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9793758392333984,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.281087276669587,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.0629,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.269,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 175000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.168345641072204,
+      "compute_loss_in_ms": 285.7807895615697,
+      "epoch": 72.76,
+      "learning_rate/full": 1.4253825239264306e-05,
+      "loss": 2.8714,
+      "step": 176000,
+      "training_step_in_ms": 1101.577396351844
+    },
+    {
+      "_prepare_inputs_in_ms": 4.185015048831701,
+      "compute_loss_in_ms": 285.48355446383357,
+      "epoch": 73.17,
+      "learning_rate/full": 1.3110908410358026e-05,
+      "loss": 2.8674,
+      "step": 177000,
+      "training_step_in_ms": 1106.790641155094
+    },
+    {
+      "_prepare_inputs_in_ms": 4.181742530316114,
+      "compute_loss_in_ms": 285.46897569671273,
+      "epoch": 73.58,
+      "learning_rate/full": 1.2015269571172228e-05,
+      "loss": 2.8685,
+      "step": 178000,
+      "training_step_in_ms": 1105.456206858158
+    },
+    {
+      "_prepare_inputs_in_ms": 4.189217183738947,
+      "compute_loss_in_ms": 285.7022790014744,
+      "epoch": 74.0,
+      "learning_rate/full": 1.0966013151343868e-05,
+      "loss": 2.8703,
+      "step": 179000,
+      "training_step_in_ms": 1103.0243016816676
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1993721053004265,
+      "compute_loss_in_ms": 285.9701578617096,
+      "epoch": 74.41,
+      "learning_rate/full": 9.962417871458617e-06,
+      "loss": 2.8679,
+      "step": 180000,
+      "training_step_in_ms": 1108.1376051008701
+    },
+    {
+      "epoch": 74.41,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.978816032409668,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.28102599122389305,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.271,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.255,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 180000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.160367134140759,
+      "compute_loss_in_ms": 285.407042812556,
+      "epoch": 74.82,
+      "learning_rate/full": 9.006738502964407e-06,
+      "loss": 2.8643,
+      "step": 181000,
+      "training_step_in_ms": 1101.3521456047893
+    },
+    {
+      "_prepare_inputs_in_ms": 4.180663630366325,
+      "compute_loss_in_ms": 285.55450112745166,
+      "epoch": 75.24,
+      "learning_rate/full": 8.098183863851083e-06,
+      "loss": 2.8683,
+      "step": 182000,
+      "training_step_in_ms": 1107.5391217172146
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1743567287921906,
+      "compute_loss_in_ms": 285.2670620009303,
+      "epoch": 75.65,
+      "learning_rate/full": 7.236978126380823e-06,
+      "loss": 2.864,
+      "step": 183000,
+      "training_step_in_ms": 1103.5778979249299
+    },
+    {
+      "_prepare_inputs_in_ms": 4.20133513212204,
+      "compute_loss_in_ms": 285.7320618443191,
+      "epoch": 76.06,
+      "learning_rate/full": 6.424124390450504e-06,
+      "loss": 2.8696,
+      "step": 184000,
+      "training_step_in_ms": 1102.3650901168585
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1793825179338455,
+      "compute_loss_in_ms": 285.5250694230199,
+      "epoch": 76.48,
+      "learning_rate/full": 5.657451579824824e-06,
+      "loss": 2.8645,
+      "step": 185000,
+      "training_step_in_ms": 1106.2330449260771
+    },
+    {
+      "epoch": 76.48,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9790234565734863,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2804143185850342,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 111.0093,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.207,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.225,
+      "step": 185000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.18695215044952,
+      "compute_loss_in_ms": 285.5290654525161,
+      "epoch": 76.89,
+      "learning_rate/full": 4.939520495279481e-06,
+      "loss": 2.8657,
+      "step": 186000,
+      "training_step_in_ms": 1101.3268077746034
+    },
+    {
+      "_prepare_inputs_in_ms": 4.185709021985531,
+      "compute_loss_in_ms": 285.5161408223212,
+      "epoch": 77.3,
+      "learning_rate/full": 4.269717665299333e-06,
+      "loss": 2.8633,
+      "step": 187000,
+      "training_step_in_ms": 1107.3937772586942
+    },
+    {
+      "_prepare_inputs_in_ms": 4.219432931393385,
+      "compute_loss_in_ms": 285.7696287557483,
+      "epoch": 77.72,
+      "learning_rate/full": 3.6476104696328672e-06,
+      "loss": 2.8611,
+      "step": 188000,
+      "training_step_in_ms": 1103.4366898052394
+    },
+    {
+      "_prepare_inputs_in_ms": 4.222444631159306,
+      "compute_loss_in_ms": 285.7456459365785,
+      "epoch": 78.13,
+      "learning_rate/full": 3.0745965927555298e-06,
+      "loss": 2.8682,
+      "step": 189000,
+      "training_step_in_ms": 1102.7803975529969
+    },
+    {
+      "_prepare_inputs_in_ms": 4.181710965931416,
+      "compute_loss_in_ms": 285.33873960748315,
+      "epoch": 78.54,
+      "learning_rate/full": 2.550171112510902e-06,
+      "loss": 2.8619,
+      "step": 190000,
+      "training_step_in_ms": 1105.0700605846941
+    },
+    {
+      "epoch": 78.54,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.979418992996216,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.28126194557201906,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 110.0926,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.267,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 190000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1674747299857255,
+      "compute_loss_in_ms": 285.1683004386723,
+      "epoch": 78.96,
+      "learning_rate/full": 2.07446342303903e-06,
+      "loss": 2.8623,
+      "step": 191000,
+      "training_step_in_ms": 1101.7088311165571
+    },
+    {
+      "_prepare_inputs_in_ms": 4.192074902355671,
+      "compute_loss_in_ms": 285.65734274312854,
+      "epoch": 79.37,
+      "learning_rate/full": 1.6475908980941423e-06,
+      "loss": 2.8592,
+      "step": 192000,
+      "training_step_in_ms": 1110.4893043078482
+    },
+    {
+      "_prepare_inputs_in_ms": 4.188501738011837,
+      "compute_loss_in_ms": 285.3913672603667,
+      "epoch": 79.79,
+      "learning_rate/full": 1.2693051031663184e-06,
+      "loss": 2.8693,
+      "step": 193000,
+      "training_step_in_ms": 1105.3064069263637
+    },
+    {
+      "_prepare_inputs_in_ms": 4.1890876069664955,
+      "compute_loss_in_ms": 285.6683066636324,
+      "epoch": 80.2,
+      "learning_rate/full": 9.404559306640304e-07,
+      "loss": 2.8679,
+      "step": 194000,
+      "training_step_in_ms": 1106.7661010883749
+    },
+    {
+      "_prepare_inputs_in_ms": 4.195614516735077,
+      "compute_loss_in_ms": 285.7047406025231,
+      "epoch": 80.61,
+      "learning_rate/full": 6.607217220689466e-07,
+      "loss": 2.8595,
+      "step": 195000,
+      "training_step_in_ms": 1110.2134825922549
+    },
+    {
+      "epoch": 80.61,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9787235260009766,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.2807572494534342,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.6903,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.293,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.228,
+      "step": 195000
+    },
+    {
+      "_prepare_inputs_in_ms": 4.209323071125077,
+      "compute_loss_in_ms": 285.70444655045867,
+      "epoch": 81.03,
+      "learning_rate/full": 4.301714976112869e-07,
+      "loss": 2.8627,
+      "step": 196000,
+      "training_step_in_ms": 1106.642473962158
+    },
+    {
+      "_prepare_inputs_in_ms": 4.216882940381765,
+      "compute_loss_in_ms": 285.90086993202567,
+      "epoch": 81.44,
+      "learning_rate/full": 2.4870533697582963e-07,
+      "loss": 2.861,
+      "step": 197000,
+      "training_step_in_ms": 1110.1682490482926
+    },
+    {
+      "_prepare_inputs_in_ms": 4.208073288202286,
+      "compute_loss_in_ms": 285.9070298522711,
+      "epoch": 81.85,
+      "learning_rate/full": 1.167309390885718e-07,
+      "loss": 2.8634,
+      "step": 198000,
+      "training_step_in_ms": 1105.8664784356952
+    },
+    {
+      "_prepare_inputs_in_ms": 4.220877002924681,
+      "compute_loss_in_ms": 285.8088936395943,
+      "epoch": 82.27,
+      "learning_rate/full": 3.407474692453949e-08,
+      "loss": 2.8649,
+      "step": 199000,
+      "training_step_in_ms": 1108.864688232541
+    },
+    {
+      "_prepare_inputs_in_ms": 4.200665093958378,
+      "compute_loss_in_ms": 285.4954933747649,
+      "epoch": 82.68,
+      "learning_rate/full": 7.485261950046507e-10,
+      "loss": 2.8661,
+      "step": 200000,
+      "training_step_in_ms": 1107.0144655555487
+    },
+    {
+      "epoch": 82.68,
+      "eval_visual_genome-densecap-local-densecap-test_loss": 2.9789507389068604,
+      "eval_visual_genome-densecap-local-densecap-test_meteor": 0.28079051577715114,
+      "eval_visual_genome-densecap-local-densecap-test_runtime": 109.9761,
+      "eval_visual_genome-densecap-local-densecap-test_samples_per_second": 7.274,
+      "eval_visual_genome-densecap-local-densecap-test_steps_per_second": 0.227,
+      "step": 200000
+    }
+  ],
+  "max_steps": 200000,
+  "num_train_epochs": 83,
+  "total_flos": 4.81380362914756e+23,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a4981a2c6d5212e3c5f9722334ed6160131dd453f8bc64be7b40cf8fa9ff105
+size 5240

zero_to_fp32.py

@@ -0,0 +1,587 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dict = torch.load(f, map_location=device)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage <= 2:
+        fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in param_shapes.items():
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
+            0).narrow(0, 0, unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file, tag=args.tag)