├── .gitignore ├── LICENSE ├── README.md ├── benchmarks ├── ResNet-50_v1.5_Performance_Comparison_TensorFlow_1.12_GCP.md └── ShapeMask_Performance_Comparison_TensorFlow_1.14_GCP.md ├── models ├── common │ ├── __init__.py │ ├── imagenet.py │ ├── inference_warmup.py │ └── tpu_profiler_hook.py ├── experimental │ ├── cifar_keras │ │ ├── README.md │ │ └── cifar_keras.py │ ├── dcgan │ │ ├── README.md │ │ ├── cifar_input.py │ │ ├── cifar_model.py │ │ ├── dcgan_main.py │ │ ├── mnist_input.py │ │ └── mnist_model.py │ ├── deeplab │ │ ├── README.md │ │ ├── data_pipeline.py │ │ ├── main.py │ │ └── model.py │ ├── densenet_keras │ │ ├── README.md │ │ ├── densenet_keras_imagenet.py │ │ ├── densenet_keras_model.py │ │ └── vgg_preprocessing.py │ ├── distribution_strategy │ │ ├── imagenet_input.py │ │ ├── resnet_estimator.py │ │ ├── resnet_model.py │ │ └── resnet_preprocessing.py │ ├── embedding │ │ ├── README.md │ │ ├── create_data.py │ │ └── model.py │ ├── inception │ │ ├── imagenet.py │ │ ├── inception_preprocessing.py │ │ ├── inception_v2.py │ │ ├── inception_v2_tpu_model.py │ │ ├── inception_v3.py │ │ ├── inception_v3_k8s.yaml │ │ ├── inception_v3_old.py │ │ ├── inception_v4.py │ │ ├── inception_v4_model.py │ │ └── vgg_preprocessing.py │ ├── inference │ │ ├── api_config.yaml │ │ ├── api_descriptor.pb │ │ ├── docker │ │ │ ├── Dockerfile-nightly.devel-tpu │ │ │ ├── Dockerfile.devel-tpu │ │ │ ├── Dockerfile.tpu │ │ │ └── README.md │ │ ├── load_test │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── data │ │ │ │ ├── __init__.py │ │ │ │ ├── criteo.py │ │ │ │ ├── data_loader.py │ │ │ │ ├── data_loader_factory.py │ │ │ │ ├── generic_jsonl.py │ │ │ │ ├── squad_bert.py │ │ │ │ ├── synthetic_bert.py │ │ │ │ ├── synthetic_bert_test.py │ │ │ │ ├── synthetic_image.py │ │ │ │ └── synthetic_image_test.py │ │ │ ├── examples │ │ │ │ ├── __init__.py │ │ │ │ ├── loadgen_grpc_main.py │ │ │ │ ├── loadgen_grpc_main_test.py │ │ │ │ ├── loadgen_rest_main.py │ │ │ │ ├── loadgen_rest_main_test.py │ │ │ │ └── loadgen_vertex_main.py │ │ │ ├── targets │ │ │ │ ├── __init__.py │ │ │ │ ├── dummy_target.py │ │ │ │ ├── grpc_target.py │ │ │ │ ├── grpc_target_test.py │ │ │ │ ├── rest_target.py │ │ │ │ ├── rest_target_test.py │ │ │ │ ├── target.py │ │ │ │ ├── target_factory.py │ │ │ │ ├── vertex_gapic_target.py │ │ │ │ ├── vertex_grpc_target.py │ │ │ │ ├── vertex_rest_target.py │ │ │ │ └── vertex_rest_target_test.py │ │ │ ├── testing │ │ │ │ ├── __init__.py │ │ │ │ └── mock_grpc.py │ │ │ ├── tools │ │ │ │ └── Dockerfile │ │ │ └── traffic_handlers │ │ │ │ ├── __init__.py │ │ │ │ ├── loadgen_handler.py │ │ │ │ ├── loadgen_handler_test.py │ │ │ │ ├── traffic_handler.py │ │ │ │ └── traffic_handler_factory.py │ │ ├── load_test_client.py │ │ ├── object_detection.ipynb │ │ ├── openapi.yaml │ │ └── setup-pool.sh │ ├── keras_application │ │ └── application_model.py │ ├── keras_colab │ │ ├── README.md │ │ └── shakespeare_lstm.py │ ├── mnist_jupyter │ │ ├── Cloud-TPU-Demo.ipynb │ │ └── launch.sh │ ├── mnist_keras │ │ ├── mnist.py │ │ └── mnist_tf2_with_summary.py │ ├── ncf │ │ ├── README.md │ │ └── ncf_main.py │ ├── resnet50_keras │ │ ├── README.md │ │ ├── imagenet_input.py │ │ ├── model_saving_utils.py │ │ ├── resnet50.py │ │ ├── resnet50_ctl_tf1.py │ │ ├── resnet50_ctl_tf2.py │ │ ├── resnet50_test.py │ │ ├── resnet50_tf2.py │ │ ├── resnet_model.py │ │ └── resnet_preprocessing.py │ └── show_and_tell │ │ ├── README │ │ ├── configuration.py │ │ ├── image_embedding.py │ │ ├── image_processing.py │ │ ├── inputs.py │ │ ├── show_and_tell_model.py │ │ └── train.py ├── hyperparameters │ ├── __init__.py │ ├── common_hparams_flags.py │ ├── common_tpu_flags.py │ ├── flags_to_params.py │ └── params_dict.py ├── official │ ├── __init__.py │ ├── amoeba_net │ │ ├── README.md │ │ ├── amoeba_net.py │ │ ├── amoeba_net_k8s.yaml │ │ ├── amoeba_net_model.py │ │ ├── inception_preprocessing.py │ │ ├── model_builder.py │ │ ├── model_specs.py │ │ ├── network_utils.py │ │ ├── network_utils_test.py │ │ └── tf_hub.py │ ├── bert │ │ └── README.md │ ├── densenet │ │ ├── README.md │ │ ├── densenet_imagenet.py │ │ ├── densenet_model.py │ │ └── vgg_preprocessing.py │ ├── detection │ │ ├── GETTING_STARTED.md │ │ ├── MODEL_ZOO.md │ │ ├── README.md │ │ ├── configs │ │ │ ├── __init__.py │ │ │ ├── base_config.py │ │ │ ├── classification_config.py │ │ │ ├── detection_config.py │ │ │ ├── factory.py │ │ │ ├── maskrcnn_config.py │ │ │ ├── retinanet_config.py │ │ │ ├── segmentation_config.py │ │ │ ├── shapemask_config.py │ │ │ ├── spinenet │ │ │ │ ├── spinenet-mbconv49_retinanet.yaml │ │ │ │ ├── spinenet143_classification.yaml │ │ │ │ ├── spinenet143_mrcnn.yaml │ │ │ │ ├── spinenet143_retinanet.yaml │ │ │ │ ├── spinenet190_classification.yaml │ │ │ │ ├── spinenet190_mrcnn.yaml │ │ │ │ ├── spinenet190_retinanet.yaml │ │ │ │ ├── spinenet49S_retinanet.yaml │ │ │ │ ├── spinenet49_classification.yaml │ │ │ │ ├── spinenet49_mrcnn.yaml │ │ │ │ ├── spinenet49_retinanet.yaml │ │ │ │ ├── spinenet49s_mrcnn.yaml │ │ │ │ ├── spinenet96_classification.yaml │ │ │ │ ├── spinenet96_mrcnn.yaml │ │ │ │ └── spinenet96_retinanet.yaml │ │ │ └── yaml │ │ │ │ ├── retinanet_autoaugment.yaml │ │ │ │ └── retinanet_nasfpn.yaml │ │ ├── dataloader │ │ │ ├── __init__.py │ │ │ ├── anchor.py │ │ │ ├── classification_parser.py │ │ │ ├── extract_objects_parser.py │ │ │ ├── factory.py │ │ │ ├── input_reader.py │ │ │ ├── input_reader_util.py │ │ │ ├── maskrcnn_parser.py │ │ │ ├── maskrcnn_parser_with_copy_paste.py │ │ │ ├── mode_keys.py │ │ │ ├── retinanet_parser.py │ │ │ ├── segmentation_parser.py │ │ │ ├── shapemask_parser.py │ │ │ └── tf_example_decoder.py │ │ ├── datasets │ │ │ └── coco_label_map.csv │ │ ├── evaluate_model.py │ │ ├── evaluation │ │ │ ├── __init__.py │ │ │ ├── coco_evaluator.py │ │ │ ├── coco_utils.py │ │ │ └── factory.py │ │ ├── executor │ │ │ ├── __init__.py │ │ │ └── tpu_executor.py │ │ ├── export_inference_graph.py │ │ ├── export_saved_model.py │ │ ├── export_tflite_model.py │ │ ├── inference.py │ │ ├── inference_saved_model.py │ │ ├── k8s │ │ │ └── retinanet_k8s.yaml │ │ ├── main.py │ │ ├── modeling │ │ │ ├── __init__.py │ │ │ ├── architecture │ │ │ │ ├── __init__.py │ │ │ │ ├── efficientnet.py │ │ │ │ ├── efficientnet_constants.py │ │ │ │ ├── factory.py │ │ │ │ ├── fpn.py │ │ │ │ ├── heads.py │ │ │ │ ├── identity.py │ │ │ │ ├── nasfpn.py │ │ │ │ ├── nn_blocks.py │ │ │ │ ├── nn_ops.py │ │ │ │ ├── resnet.py │ │ │ │ ├── spinenet.py │ │ │ │ └── spinenet_mbconv.py │ │ │ ├── base_model.py │ │ │ ├── cascade_maskrcnn_model.py │ │ │ ├── classification_model.py │ │ │ ├── factory.py │ │ │ ├── learning_rates.py │ │ │ ├── losses.py │ │ │ ├── maskrcnn_model.py │ │ │ ├── model_builder.py │ │ │ ├── retinanet_model.py │ │ │ ├── segmentation_model.py │ │ │ └── shapemask_model.py │ │ ├── ops │ │ │ ├── __init__.py │ │ │ ├── nms.py │ │ │ ├── postprocess_ops.py │ │ │ ├── roi_ops.py │ │ │ ├── spatial_transform_ops.py │ │ │ └── target_ops.py │ │ ├── projects │ │ │ ├── copy_paste │ │ │ │ ├── README.md │ │ │ │ └── configs │ │ │ │ │ ├── cascade_maskrcnn_effb7_1280.yaml │ │ │ │ │ ├── cascade_maskrcnn_effb7_nasfpn_1280.yaml │ │ │ │ │ ├── cascade_maskrcnn_effb7_nasfpn_vanilla_1280.yaml │ │ │ │ │ ├── maskrcnn_effb7_1024.yaml │ │ │ │ │ ├── maskrcnn_effb7_1280.yaml │ │ │ │ │ ├── maskrcnn_effb7_640.yaml │ │ │ │ │ ├── maskrcnn_resnet101_1024.yaml │ │ │ │ │ ├── maskrcnn_resnet101_1280.yaml │ │ │ │ │ └── maskrcnn_resnet50_1024.yaml │ │ │ ├── fashionpedia │ │ │ │ ├── README.md │ │ │ │ ├── configs │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── factory.py │ │ │ │ │ ├── model_config.py │ │ │ │ │ └── yaml │ │ │ │ │ │ ├── r101fpn_amrcnn.yaml │ │ │ │ │ │ ├── r50fpn_amrcnn.yaml │ │ │ │ │ │ ├── spinenet143_amrcnn.yaml │ │ │ │ │ │ ├── spinenet49_amrcnn.yaml │ │ │ │ │ │ └── spinenet96_amrcnn.yaml │ │ │ │ ├── dataloader │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── data_parser.py │ │ │ │ │ ├── factory.py │ │ │ │ │ ├── input_reader.py │ │ │ │ │ └── tf_example_decoder.py │ │ │ │ ├── dataset │ │ │ │ │ └── fashionpedia_label_map.csv │ │ │ │ ├── export_saved_model.py │ │ │ │ ├── inference.py │ │ │ │ ├── inference_saved_model.py │ │ │ │ ├── main.py │ │ │ │ ├── modeling │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── architecture │ │ │ │ │ │ ├── __init__.py │ │ │ │ │ │ ├── factory.py │ │ │ │ │ │ └── fast_rcnn_head.py │ │ │ │ │ ├── attribute_loss.py │ │ │ │ │ ├── factory.py │ │ │ │ │ ├── model.py │ │ │ │ │ └── model_builder.py │ │ │ │ ├── ops │ │ │ │ │ ├── __init__.py │ │ │ │ │ ├── postprocess_ops.py │ │ │ │ │ └── roi_sampler.py │ │ │ │ ├── serving │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── serving.py │ │ │ │ └── utils │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── dataloader_utils.py │ │ │ ├── openseg │ │ │ │ ├── OpenSeg_demo.ipynb │ │ │ │ └── README.md │ │ │ ├── self_training │ │ │ │ ├── README.md │ │ │ │ └── configs │ │ │ │ │ ├── coco_spinenet143_retinanet.yaml │ │ │ │ │ ├── coco_spinenet190_retinanet.yaml │ │ │ │ │ ├── pascal_seg_efficientnet-b7-nasfpn.yaml │ │ │ │ │ └── pascal_seg_efficientnet-l2-nasfpn.yaml │ │ │ └── vild │ │ │ │ ├── README.md │ │ │ │ ├── ViLD_demo.ipynb │ │ │ │ ├── configs │ │ │ │ ├── vild_config.py │ │ │ │ ├── vild_ensemble_resnet.yaml │ │ │ │ └── vild_resnet.yaml │ │ │ │ ├── dataloader │ │ │ │ ├── tf_example_decoder.py │ │ │ │ └── vild_parser.py │ │ │ │ ├── modeling │ │ │ │ ├── vild_head.py │ │ │ │ ├── vild_losses.py │ │ │ │ └── vild_model.py │ │ │ │ ├── ops │ │ │ │ └── postprocess_ops.py │ │ │ │ ├── preprocessing │ │ │ │ ├── __init__.py │ │ │ │ ├── create_lvis_tf_record.py │ │ │ │ └── dataset_util.py │ │ │ │ └── requirements.txt │ │ ├── serving │ │ │ ├── __init__.py │ │ │ ├── detection.py │ │ │ ├── inputs.py │ │ │ └── segmentation.py │ │ └── utils │ │ │ ├── __init__.py │ │ │ ├── autoaugment_utils.py │ │ │ ├── benchmark_utils.py │ │ │ ├── box_utils.py │ │ │ ├── class_utils.py │ │ │ ├── config_utils.py │ │ │ ├── dataloader_utils.py │ │ │ ├── input_utils.py │ │ │ ├── mask_utils.py │ │ │ ├── object_detection │ │ │ ├── __init__.py │ │ │ ├── argmax_matcher.py │ │ │ ├── balanced_positive_negative_sampler.py │ │ │ ├── box_coder.py │ │ │ ├── box_list.py │ │ │ ├── box_list_ops.py │ │ │ ├── faster_rcnn_box_coder.py │ │ │ ├── matcher.py │ │ │ ├── minibatch_sampler.py │ │ │ ├── ops.py │ │ │ ├── preprocessor.py │ │ │ ├── region_similarity_calculator.py │ │ │ ├── shape_utils.py │ │ │ ├── target_assigner.py │ │ │ └── visualization_utils.py │ │ │ ├── saved_model_benchmark.py │ │ │ └── shape_utils.py │ ├── efficientnet │ │ ├── README.md │ │ ├── __init__.py │ │ ├── autoaugment.py │ │ ├── condconv │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── condconv_layers.py │ │ │ └── efficientnet_condconv_builder.py │ │ ├── edgetpu │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ └── efficientnet_edgetpu_builder.py │ │ ├── efficientnet_builder.py │ │ ├── efficientnet_model.py │ │ ├── eval_ckpt_example.ipynb │ │ ├── eval_ckpt_main.py │ │ ├── export_model.py │ │ ├── g3doc │ │ │ ├── condconv-layer.png │ │ │ ├── efficientnet-edgetpu.png │ │ │ ├── flops.png │ │ │ ├── lite-float-gpu.png │ │ │ ├── lite-quant-cpu.png │ │ │ ├── lite-quant-size.png │ │ │ └── params.png │ │ ├── imagenet_input.py │ │ ├── inspect_model_architecture.py │ │ ├── lars_optimizer.py │ │ ├── lite │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── efficientnet_lite_builder.py │ │ │ ├── efficientnet_lite_builder_test.py │ │ │ ├── efficientnet_lite_model_qat.py │ │ │ └── efficientnet_lite_model_qat_test.py │ │ ├── main.py │ │ ├── model_builder_factory.py │ │ ├── preprocessing.py │ │ ├── tpu │ │ │ ├── __init__.py │ │ │ └── efficientnet_x_builder.py │ │ └── utils.py │ ├── mask_rcnn │ │ ├── README.md │ │ ├── anchors.py │ │ ├── box_utils.py │ │ ├── coco_metric.py │ │ ├── coco_utils.py │ │ ├── configs │ │ │ ├── __init__.py │ │ │ ├── cloud │ │ │ │ ├── v2-128.yaml │ │ │ │ ├── v2-32.yaml │ │ │ │ ├── v2-8.yaml │ │ │ │ ├── v3-128.yaml │ │ │ │ ├── v3-32.yaml │ │ │ │ └── v3-8.yaml │ │ │ └── mask_rcnn_config.py │ │ ├── dataloader.py │ │ ├── distributed_executer.py │ │ ├── evaluation.py │ │ ├── export_saved_model.py │ │ ├── fpn.py │ │ ├── heads.py │ │ ├── learning_rates.py │ │ ├── losses.py │ │ ├── mask_rcnn_demo.ipynb │ │ ├── mask_rcnn_k8s.yaml │ │ ├── mask_rcnn_main.py │ │ ├── mask_rcnn_model.py │ │ ├── nn_ops.py │ │ ├── object_detection │ │ │ ├── __init__.py │ │ │ ├── argmax_matcher.py │ │ │ ├── balanced_positive_negative_sampler.py │ │ │ ├── box_coder.py │ │ │ ├── box_list.py │ │ │ ├── faster_rcnn_box_coder.py │ │ │ ├── matcher.py │ │ │ ├── minibatch_sampler.py │ │ │ ├── ops.py │ │ │ ├── preprocessor.py │ │ │ ├── region_similarity_calculator.py │ │ │ ├── shape_utils.py │ │ │ ├── target_assigner.py │ │ │ ├── tf_example_decoder.py │ │ │ └── visualization_utils.py │ │ ├── postprocess_ops.py │ │ ├── preprocess_ops.py │ │ ├── resnet.py │ │ ├── roi_ops.py │ │ ├── serving.py │ │ ├── spatial_transform_ops.py │ │ ├── tpu_normalization.py │ │ └── training_ops.py │ ├── mnasnet │ │ ├── README.md │ │ ├── configs │ │ │ ├── __init__.py │ │ │ ├── cloud │ │ │ │ ├── gpu.yaml │ │ │ │ ├── v2-32.yaml │ │ │ │ ├── v2-8.yaml │ │ │ │ ├── v3-32.yaml │ │ │ │ └── v3-8.yaml │ │ │ └── mnasnet_config.py │ │ ├── eval_ckpt_main.py │ │ ├── g3doc │ │ │ ├── mnasnet_vs_mobilenetv2.png │ │ │ └── mnasnet_vs_mobilenetv2_2.png │ │ ├── imagenet_input.py │ │ ├── mixnet │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── custom_layers.py │ │ │ ├── g3doc │ │ │ │ └── mixnet-flops.png │ │ │ ├── mixnet_builder.py │ │ │ ├── mixnet_eval_example.ipynb │ │ │ └── mixnet_model.py │ │ ├── mnas_utils.py │ │ ├── mnasnet_example.ipynb │ │ ├── mnasnet_main.py │ │ ├── mnasnet_model.py │ │ ├── mnasnet_models.py │ │ ├── post_quantization.py │ │ └── preprocessing.py │ ├── mnist │ │ ├── README.md │ │ └── mnist_tpu.py │ ├── mobilenet │ │ ├── README.md │ │ ├── configs │ │ │ ├── __init__.py │ │ │ ├── cloud │ │ │ │ ├── v2-128.yaml │ │ │ │ ├── v2-256.yaml │ │ │ │ ├── v2-32.yaml │ │ │ │ ├── v2-512.yaml │ │ │ │ ├── v2-8.yaml │ │ │ │ ├── v3-1024.yaml │ │ │ │ ├── v3-128.yaml │ │ │ │ ├── v3-2048.yaml │ │ │ │ ├── v3-256.yaml │ │ │ │ ├── v3-32.yaml │ │ │ │ ├── v3-512.yaml │ │ │ │ ├── v3-64.yaml │ │ │ │ └── v3-8.yaml │ │ │ └── mobilenet_config.py │ │ ├── inception_preprocessing.py │ │ ├── mobilenet.py │ │ ├── mobilenet_model.py │ │ ├── supervised_images.py │ │ └── vgg_preprocessing.py │ ├── resnet │ │ ├── README.md │ │ ├── __init__.py │ │ ├── benchmark │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── read_training_time.py │ │ │ └── resnet_benchmark.py │ │ ├── configs │ │ │ ├── __init__.py │ │ │ ├── cloud │ │ │ │ ├── randaugment-32.yaml │ │ │ │ ├── randaugment-8.yaml │ │ │ │ ├── v2-128.yaml │ │ │ │ ├── v2-256.yaml │ │ │ │ ├── v2-32.yaml │ │ │ │ ├── v2-512.yaml │ │ │ │ ├── v2-8.yaml │ │ │ │ ├── v3-1024.yaml │ │ │ │ ├── v3-128.yaml │ │ │ │ ├── v3-2048.yaml │ │ │ │ ├── v3-256.yaml │ │ │ │ ├── v3-32.yaml │ │ │ │ ├── v3-512.yaml │ │ │ │ ├── v3-64.yaml │ │ │ │ └── v3-8.yaml │ │ │ └── resnet_config.py │ │ ├── imagenet_input.py │ │ ├── lars_util.py │ │ ├── resnet_k8s.yaml │ │ ├── resnet_layers.py │ │ ├── resnet_main.py │ │ ├── resnet_model.py │ │ ├── resnet_model_test.py │ │ ├── resnet_preprocessing.py │ │ └── resnet_rs │ │ │ ├── README.md │ │ │ └── configs │ │ │ ├── resnetrs101_i160.yaml │ │ │ ├── resnetrs101_i192.yaml │ │ │ ├── resnetrs152_i192.yaml │ │ │ ├── resnetrs152_i224.yaml │ │ │ ├── resnetrs152_i256.yaml │ │ │ ├── resnetrs200_i256.yaml │ │ │ ├── resnetrs270_i256.yaml │ │ │ ├── resnetrs350_i256.yaml │ │ │ ├── resnetrs350_i320.yaml │ │ │ ├── resnetrs420_i320.yaml │ │ │ └── resnetrs50_i160.yaml │ ├── retinanet │ │ ├── README.md │ │ ├── anchors.py │ │ ├── coco_metric.py │ │ ├── dataloader.py │ │ ├── evaluation.py │ │ ├── object_detection │ │ │ ├── __init__.py │ │ │ ├── argmax_matcher.py │ │ │ ├── box_coder.py │ │ │ ├── box_list.py │ │ │ ├── faster_rcnn_box_coder.py │ │ │ ├── matcher.py │ │ │ ├── preprocessor.py │ │ │ ├── region_similarity_calculator.py │ │ │ ├── shape_utils.py │ │ │ ├── target_assigner.py │ │ │ └── tf_example_decoder.py │ │ ├── postprocess.py │ │ ├── retinanet_architecture.py │ │ ├── retinanet_k8s.yaml │ │ ├── retinanet_main.py │ │ ├── retinanet_model.py │ │ ├── retinanet_segmentation_main.py │ │ ├── retinanet_segmentation_model.py │ │ └── retinanet_tensorrt.py │ ├── squeezenet │ │ ├── configs │ │ │ ├── __init__.py │ │ │ └── squeezenet_config.py │ │ ├── data_pipeline.py │ │ ├── squeezenet_main.py │ │ └── squeezenet_model.py │ ├── transformer │ │ └── README.md │ └── unet3d │ │ ├── README.md │ │ ├── configs │ │ └── cloud │ │ │ ├── v3-128_256x256x256_ce.yaml │ │ │ ├── v3-128_256x256x256_dice.yaml │ │ │ ├── v3-32_256x256x256_ce.yaml │ │ │ ├── v3-32_256x256x256_dice.yaml │ │ │ ├── v3-8_128x128x128_ce.yaml │ │ │ └── v3-8_128x128x128_dice.yaml │ │ ├── data_preprocess │ │ ├── convert_lits.py │ │ ├── convert_lits_nii_to_npy.py │ │ └── generate_fake_lits.py │ │ ├── export_saved_model.py │ │ ├── input_reader.py │ │ ├── metrics.py │ │ ├── requirements.txt │ │ ├── saved_model_inference.py │ │ ├── tpu_executor.py │ │ ├── unet_config.py │ │ ├── unet_main.py │ │ └── unet_model.py ├── samples │ └── core │ │ └── get_started │ │ ├── custom_tpuestimator.py │ │ └── iris_data_tpu.py └── setup.py └── tools ├── colab ├── BUILD ├── autoencoder_embeddings.ipynb ├── bert_finetuning_with_cloud_tpus.ipynb ├── classification_iris_data_with_keras.ipynb ├── classification_iris_data_with_tpuestimator.ipynb ├── custom_training.ipynb ├── fashion_mnist.ipynb ├── image_classification_converter.ipynb ├── keras_mnist_tpu.ipynb ├── mnist_estimator.ipynb ├── mnist_tpuestimator.ipynb ├── profiling_tpus_in_colab.ipynb ├── regression_sine_data_with_keras.ipynb ├── shakespeare_with_tpu_and_keras.ipynb ├── shakespeare_with_tpuestimator.ipynb ├── shapemask_instance_segmentation.ipynb └── tfgan_on_tpus.ipynb ├── ctpu ├── .gitignore ├── README.md ├── commands │ ├── auth.go │ ├── auth_test.go │ ├── common.go │ ├── common_test.go │ ├── config_cmd.go │ ├── delete.go │ ├── delete_test.go │ ├── list.go │ ├── list_test.go │ ├── pause.go │ ├── pause_test.go │ ├── quota.go │ ├── restart.go │ ├── restart_test.go │ ├── status.go │ ├── status_test.go │ ├── tf_versions.go │ ├── tpu_locations.go │ ├── tpu_locations_test.go │ ├── tpu_size.go │ ├── tpu_size_test.go │ ├── up.go │ ├── up_test.go │ └── version.go ├── config │ ├── config.go │ ├── config_gce.go │ ├── config_gcloud.go │ ├── config_gcloud_test.go │ ├── config_test.go │ ├── devshell.go │ ├── devshell_test.go │ └── testdata │ │ └── gcloud │ │ ├── clean │ │ ├── README.md │ │ ├── active_config │ │ ├── application_default_credentials.json │ │ └── configurations │ │ │ └── config_ctpu9 │ │ ├── corrupted │ │ ├── README.md │ │ ├── active_config │ │ ├── application_default_credentials.json │ │ └── configurations │ │ │ └── config_default │ │ ├── corrupted2 │ │ ├── README.md │ │ ├── active_config │ │ └── application_default_credentials.json │ │ ├── incomplete │ │ ├── README.md │ │ ├── active_config │ │ ├── application_default_credentials.json │ │ └── configurations │ │ │ └── config_ctpu9 │ │ ├── no_app_creds │ │ ├── README.md │ │ ├── active_config │ │ └── configurations │ │ │ └── config_ctpu9 │ │ └── no_config │ │ └── README.md ├── ctrl │ ├── ctrl.go │ ├── ctrl_test.go │ ├── gce.go │ ├── gce_test.go │ ├── gcloud_cli.go │ ├── gcloud_cli_test.go │ ├── resourcemgmt.go │ ├── resourcemgmt_test.go │ ├── serviceusage.go │ ├── tpu.go │ └── tpu_test.go ├── main.go └── tutorial.md ├── data_converter ├── README.md ├── __init__.py ├── image_classification │ ├── __init__.py │ ├── image_classification_data.py │ └── simple_example.py ├── image_utils.py └── object_detection │ ├── __init__.py │ ├── bbox_utils.py │ ├── object_detection_data.py │ └── simple_example.py ├── dataset_profiler ├── __init__.py └── profile_imagenet.py ├── datasets ├── README.md ├── __init__.py ├── create_coco_tf_record.py ├── download_and_preprocess_coco.sh ├── download_and_preprocess_coco_k8s.yaml ├── imagenet_to_gcs.py ├── imagenet_to_gcs_k8s.yaml └── jpeg_to_tf_record.py ├── diagnostics └── diagnostics.py ├── docker ├── Dockerfile └── Dockerfile.util ├── driver ├── Makefile.tpu_kernel_modules ├── README ├── drivers │ ├── accel │ │ ├── Kconfig │ │ ├── Makefile │ │ ├── accel.c │ │ └── accel_lib.c │ ├── asic_sw │ │ ├── Kconfig │ │ ├── Makefile │ │ ├── asic_fw_device_owner_accessor.h │ │ ├── asic_fw_device_owner_offsets.h │ │ ├── asic_fw_host_reset_offsets.h │ │ ├── asic_fw_indirect_accessor.h │ │ ├── asic_fw_indirect_register_offsets.h │ │ ├── asic_fw_reinit_reset_offsets.h │ │ ├── asic_fw_state_offsets.h │ │ ├── asic_fw_version_offsets.h │ │ ├── asic_sw_clock.c │ │ ├── asic_sw_clock.h │ │ ├── asic_sw_firmware_indirect_registers.c │ │ ├── asic_sw_firmware_indirect_registers.h │ │ └── asic_sw_module.c │ ├── char │ │ └── tpu_common │ │ │ ├── Kconfig │ │ │ ├── Makefile │ │ │ ├── tpu_common.c │ │ │ ├── tpu_common.h │ │ │ ├── tpuv2_core.c │ │ │ ├── tpuv2_core.h │ │ │ ├── tpuv2_driver.c │ │ │ ├── tpuv3_driver.c │ │ │ ├── tpuv4.c │ │ │ ├── tpuv4.h │ │ │ ├── tpuv4 │ │ │ ├── tpuv4.h │ │ │ ├── tpuv4_bar0_ranges.c │ │ │ ├── tpuv4_bar0_ranges.h │ │ │ ├── tpuv4_bar2_ranges.c │ │ │ ├── tpuv4_bar2_ranges.h │ │ │ ├── tpuv4_device_owner_offsets.c │ │ │ ├── tpuv4_device_owner_offsets.h │ │ │ ├── tpuv4_firmware_state_offsets.c │ │ │ ├── tpuv4_firmware_state_offsets.h │ │ │ ├── tpuv4_firmware_version_offsets.c │ │ │ ├── tpuv4_firmware_version_offsets.h │ │ │ ├── tpuv4_interrupt_desc.h │ │ │ ├── tpuv4_mgr_error_loperf_offsets.c │ │ │ ├── tpuv4_mgr_error_loperf_offsets.h │ │ │ ├── tpuv4_page_table_config.h │ │ │ ├── tpuv4_pcie_flr_status_offsets.c │ │ │ ├── tpuv4_pcie_flr_status_offsets.h │ │ │ ├── tpuv4_reinit_reset_offsets.c │ │ │ ├── tpuv4_reinit_reset_offsets.h │ │ │ ├── tpuv4_userspace_firmware_indirect_accessor_offsets.c │ │ │ ├── tpuv4_userspace_firmware_indirect_accessor_offsets.h │ │ │ ├── tpuv4_userspace_lst_port_indirect_offsets.c │ │ │ └── tpuv4_userspace_lst_port_indirect_offsets.h │ │ │ ├── tpuv4_driver.c │ │ │ ├── tpuv4common.c │ │ │ ├── tpuv4common.h │ │ │ ├── tpuv4common │ │ │ ├── tpuv4common_lst.c │ │ │ ├── tpuv4common_lst.h │ │ │ ├── tpuv4common_mgr_error_loperf_offsets.h │ │ │ ├── tpuv4common_pcie_flr_status_offsets.h │ │ │ ├── tpuv4common_userspace_lst_port_indirect_accessor.h │ │ │ └── tpuv4common_userspace_lst_port_indirect_offsets.h │ │ │ ├── tpuv4lite │ │ │ ├── tpuv4lite.h │ │ │ ├── tpuv4lite_bar0_ranges.c │ │ │ ├── tpuv4lite_bar0_ranges.h │ │ │ ├── tpuv4lite_bar2_ranges.c │ │ │ ├── tpuv4lite_bar2_ranges.h │ │ │ ├── tpuv4lite_device_owner_offsets.c │ │ │ ├── tpuv4lite_device_owner_offsets.h │ │ │ ├── tpuv4lite_firmware_state_offsets.c │ │ │ ├── tpuv4lite_firmware_state_offsets.h │ │ │ ├── tpuv4lite_firmware_version_offsets.c │ │ │ ├── tpuv4lite_firmware_version_offsets.h │ │ │ ├── tpuv4lite_interrupt_desc.h │ │ │ ├── tpuv4lite_mgr_error_loperf_offsets.c │ │ │ ├── tpuv4lite_mgr_error_loperf_offsets.h │ │ │ ├── tpuv4lite_page_table_config.h │ │ │ ├── tpuv4lite_pcie_flr_status_offsets.c │ │ │ ├── tpuv4lite_pcie_flr_status_offsets.h │ │ │ ├── tpuv4lite_reinit_reset_offsets.c │ │ │ ├── tpuv4lite_reinit_reset_offsets.h │ │ │ ├── tpuv4lite_userspace_firmware_indirect_accessor_offsets.c │ │ │ ├── tpuv4lite_userspace_firmware_indirect_accessor_offsets.h │ │ │ ├── tpuv4lite_userspace_lst_port_indirect_offsets.c │ │ │ └── tpuv4lite_userspace_lst_port_indirect_offsets.h │ │ │ └── tpuv4lite_driver.c │ └── gasket │ │ ├── Kconfig │ │ ├── Makefile │ │ ├── README │ │ ├── gasket_constants.h │ │ ├── gasket_core.c │ │ ├── gasket_core.h │ │ ├── gasket_dmabuf.c │ │ ├── gasket_dmabuf.h │ │ ├── gasket_interrupt.c │ │ ├── gasket_interrupt.h │ │ ├── gasket_ioctl.c │ │ ├── gasket_ioctl.h │ │ ├── gasket_logging.h │ │ ├── gasket_page_table.c │ │ ├── gasket_page_table.h │ │ ├── gasket_sysfs.c │ │ ├── gasket_sysfs.h │ │ └── gasket_types.h └── include │ ├── linux │ ├── accel.h │ ├── accel_lib.h │ ├── google │ │ ├── asic_fw_device_owner_accessor.h │ │ └── gasket.h │ └── tpuv2_ioctl.h │ ├── trace │ └── events │ │ ├── gasket_interrupt.h │ │ ├── gasket_ioctl.h │ │ └── gasket_mmap.h │ └── uapi │ └── linux │ └── accel_ioctl.h ├── grpc_tpu_worker └── README.md ├── kubernetes ├── tensorboard_k8s.yaml ├── terraform │ ├── batching_with_compact_placement │ │ ├── add_node_pool │ │ │ ├── examples │ │ │ │ ├── v4 │ │ │ │ │ ├── main.tf │ │ │ │ │ ├── outputs.tf │ │ │ │ │ └── terraform.tfvars │ │ │ │ └── v5e │ │ │ │ │ ├── main.tf │ │ │ │ │ ├── outputs.tf │ │ │ │ │ └── terraform.tfvars │ │ │ └── module │ │ │ │ ├── main.tf │ │ │ │ ├── outputs.tf │ │ │ │ ├── terraform.tfvars │ │ │ │ └── variables.tf │ │ └── create_cluster │ │ │ ├── examples │ │ │ ├── v4 │ │ │ │ ├── main.tf │ │ │ │ ├── outputs.tf │ │ │ │ └── terraform.tfvars │ │ │ └── v5e │ │ │ │ ├── main.tf │ │ │ │ ├── outputs.tf │ │ │ │ └── terraform.tfvars │ │ │ └── module │ │ │ ├── main.tf │ │ │ ├── outputs.tf │ │ │ ├── terraform.tfvars │ │ │ └── variables.tf │ ├── batching_without_compact_placement │ │ ├── add_node_pool │ │ │ ├── examples │ │ │ │ ├── v4 │ │ │ │ │ ├── main.tf │ │ │ │ │ ├── outputs.tf │ │ │ │ │ └── terraform.tfvars │ │ │ │ └── v5e │ │ │ │ │ ├── main.tf │ │ │ │ │ ├── outputs.tf │ │ │ │ │ └── terraform.tfvars │ │ │ └── module │ │ │ │ ├── main.tf │ │ │ │ ├── outputs.tf │ │ │ │ ├── terraform.tfvars │ │ │ │ └── variables.tf │ │ └── create_cluster │ │ │ ├── examples │ │ │ ├── v4 │ │ │ │ ├── main.tf │ │ │ │ ├── outputs.tf │ │ │ │ └── terraform.tfvars │ │ │ └── v5e │ │ │ │ ├── main.tf │ │ │ │ ├── outputs.tf │ │ │ │ └── terraform.tfvars │ │ │ └── module │ │ │ ├── main.tf │ │ │ ├── outputs.tf │ │ │ ├── terraform.tfvars │ │ │ └── variables.tf │ └── non_batching_with_compact_placement │ │ ├── examples │ │ ├── v4 │ │ │ ├── main.tf │ │ │ ├── outputs.tf │ │ │ └── terraform.tfvars │ │ └── v5e │ │ │ ├── main.tf │ │ │ ├── outputs.tf │ │ │ └── terraform.tfvars │ │ └── module │ │ ├── main.tf │ │ ├── outputs.tf │ │ ├── terraform.tfvars │ │ └── variables.tf └── tpu_profiler_k8s.yaml ├── ray_tpu ├── README.md ├── create_cpu.sh ├── create_tpu_service_account.sh ├── deploy.sh ├── legacy │ ├── README.md │ ├── ipp_tool.py │ ├── jax_example.ipynb │ ├── jupyter_multihost_tpu.md │ ├── ray_tpu_controller.py │ ├── requirements-notebook.txt │ ├── run_basic_jax.py │ ├── run_pax_autoresume.py │ ├── tpu_api.py │ ├── tpu_controller.py │ └── tpu_ray_cluster.png └── src │ ├── requirements.txt │ ├── serve │ ├── README.md │ ├── cluster.yaml │ ├── fake_load_test.py │ ├── ray_serve_diffusion_flax.py │ ├── requirements.txt │ └── set_ray_address.sh │ └── tune │ ├── README.md │ ├── cluster.yaml │ ├── requirements.txt │ ├── run_hp_search.py │ └── set_ray_address.sh └── retry ├── README.md └── retry.sh /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | .ipynb_checkpoints 3 | node_modules 4 | /.bazelrc 5 | /.tf_configure.bazelrc 6 | /bazel-* 7 | /bazel_pip 8 | /third_party/eigen3/mkl_include 9 | /third_party/mkl/* 10 | /tools/python_bin_path.sh 11 | /tools/git/gen 12 | /pip_test 13 | /_python_build 14 | *.pyc 15 | __pycache__ 16 | *.swp 17 | .vscode/ -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Cloud TPUs # 2 | 3 | This repository is a collection of reference models and tools used with 4 | [Cloud TPUs](https://cloud.google.com/tpu/). 5 | 6 | The fastest way to get started training a model on a Cloud TPU is by following 7 | the tutorial. Click the button below to launch the tutorial using Google Cloud 8 | Shell. 9 | 10 | [![Open in Cloud Shell](http://gstatic.com/cloudssh/images/open-btn.svg)](https://console.cloud.google.com/cloudshell/open?git_repo=https%3A%2F%2Fgithub.com%2Ftensorflow%2Ftpu&page=shell&tutorial=tools%2Fctpu%2Ftutorial.md) 11 | 12 | _Note:_ This repository is a public mirror, pull requests will not be accepted. 13 | Please file an issue if you have a feature or bug request. 14 | 15 | ## Running Models 16 | 17 | To run models in the `models` subdirectory, you may need to add the top-level 18 | `/models` folder to the Python path with the command: 19 | 20 | ``` 21 | export PYTHONPATH="$PYTHONPATH:/path/to/models" 22 | ``` 23 | -------------------------------------------------------------------------------- /models/common/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/experimental/cifar_keras/README.md: -------------------------------------------------------------------------------- 1 | # Cifar Keras # 2 | 3 | This directory contains an example using the Keras layers API inside an 4 | Estimator/TPUEstimator. If you have a complete Keras model already built, 5 | consider the new experimental Cloud TPU-Keras integration available since TF 6 | 1.9. For examples, see [`models/experimental/keras`](https://github.com/tensorflow/tpu/tree/master/models/experimental/keras) 7 | -------------------------------------------------------------------------------- /models/experimental/dcgan/README.md: -------------------------------------------------------------------------------- 1 | ## Overview 2 | 3 | This example uses a DCGAN architecture to learn to produce MNIST digits and 4 | CIFAR10 images. It trains on Google Cloud TPUs. It uses an open source library 5 | called TF-GAN to abstract away many of the GAN and TPU infrastructure details. 6 | 7 | To run this example, be sure to install TF-GAN with: 8 | 9 | pip install tensorflow-gan 10 | -------------------------------------------------------------------------------- /models/experimental/densenet_keras/README.md: -------------------------------------------------------------------------------- 1 | # Cloud TPU Port of DenseNet using Keras API 2 | 3 | This folder contains an implementation of the [DenseNet](https://arxiv.org/pdf/1608.06993.pdf) 4 | image classification model using the Keras API. 5 | 6 | ## Running the model on ImageNet 7 | 8 | The process for running on ImageNet is similar, just specify the directory 9 | containing your converted tfrecord files: 10 | 11 | ``` 12 | python densenet_keras_imagenet.py\ 13 | --alsologtostderr\ 14 | --num_shards=8\ 15 | --batch_size=1024\ 16 | --master=grpc://$TPU_WORKER:8470\ 17 | --use_tpu=1\ 18 | --model_dir=gs://my-cloud-bucket/models/densenet-keras-imagenet/0\ 19 | --data_dir=gs://my-cloud-bucket/data/imagenet 20 | -------------------------------------------------------------------------------- /models/experimental/inference/api_descriptor.pb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/experimental/inference/api_descriptor.pb -------------------------------------------------------------------------------- /models/experimental/inference/load_test/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/experimental/inference/load_test/data/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/experimental/inference/load_test/examples/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/experimental/inference/load_test/targets/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/experimental/inference/load_test/testing/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/experimental/inference/load_test/traffic_handlers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/experimental/inference/load_test/traffic_handlers/traffic_handler_factory.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Target factory.""" 16 | 17 | from absl import logging 18 | 19 | from load_test.traffic_handlers import loadgen_handler 20 | from load_test.traffic_handlers import traffic_handler 21 | 22 | 23 | def get_traffic_handler( 24 | name: str, **kwargs) -> traffic_handler.TrafficHandler: 25 | if name == "loadgen": 26 | logging.info("Creating loadgen traffic handler.") 27 | return loadgen_handler.LoadGenHandler(**kwargs) 28 | else: 29 | raise ValueError("Unsupported traffic handler.") 30 | -------------------------------------------------------------------------------- /models/experimental/keras_colab/README.md: -------------------------------------------------------------------------------- 1 | # Keras Colab # 2 | 3 | This directory contains an example of using the experimental Cloud TPU-Keras 4 | integration that was added in TF 1.9 in an interactive collaboratory 5 | environment. To learn more about this new integration, 6 | check out the documentation (coming soon!). 7 | -------------------------------------------------------------------------------- /models/experimental/resnet50_keras/README.md: -------------------------------------------------------------------------------- 1 | # ResNet-50 # 2 | 3 | This directory contains an example of using the experimental Cloud TPU-Keras 4 | integration that was added in TF 1.9. ResNet-50 is a commonly used convolutional 5 | neural network used for image classification. The ResNet-family of models were 6 | introduced for the ImageNet 2015 competition and performed very well there. 7 | -------------------------------------------------------------------------------- /models/experimental/show_and_tell/README: -------------------------------------------------------------------------------- 1 | This model demonstrates using TPUEstimator to train the ShowAndTell image captioning model. 2 | 3 | The original model source is from the Tensorflow models repository: 4 | 5 | https://github.com/tensorflow/models/blob/master/research/im2txt/ 6 | 7 | You will need TFRecords from the COCO dataset for training. These can be generated by following 8 | the directions here: 9 | 10 | https://cloud.google.com/tpu/docs/tutorials/retinanet#prepare_the_coco_dataset 11 | 12 | To run the trainer: 13 | 14 | python train.py\ 15 | --tpu=$TPU_WORKER \ 16 | --input_file_pattern=gs://bucket/mscoco-records/train-*\ 17 | --model_dir=$MODEL_DIR\ 18 | --train_steps=10000\ 19 | --train_batch_size=1024 20 | 21 | 22 | NOTE: 23 | 24 | This version uses the non-reversible tf.strings.to_hash_bucket to hash captions to indices: it is 25 | not usable for inference! -------------------------------------------------------------------------------- /models/hyperparameters/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/__init__.py -------------------------------------------------------------------------------- /models/official/bert/README.md: -------------------------------------------------------------------------------- 1 | See https://github.com/google-research/bert/blob/master/README.md 2 | -------------------------------------------------------------------------------- /models/official/densenet/README.md: -------------------------------------------------------------------------------- 1 | # Cloud TPU Port of DenseNet 2 | 3 | This folder contains an implementation of the [DenseNet](https://arxiv.org/pdf/1608.06993.pdf) 4 | image classification model. 5 | 6 | ## Running the model on ImageNet 7 | 8 | The process for running on ImageNet is similar, just specify the directory 9 | containing your converted tfrecord files: 10 | 11 | ``` 12 | python densenet_imagenet.py\ 13 | --alsologtostderr\ 14 | --num_shards=8\ 15 | --batch_size=1024\ 16 | --master=grpc://$TPU_WORKER:8470\ 17 | --use_tpu=1\ 18 | --model_dir=gs://my-cloud-bucket/models/densenet-imagenet/0\ 19 | --data_dir=gs://my-cloud-bucket/data/imagenet 20 | -------------------------------------------------------------------------------- /models/official/detection/configs/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet-mbconv49_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet49-MBConv + RetinaNet with swish. 28.5% mAP. 2 | architecture: 3 | backbone: 'spinenet_mbconv' 4 | multilevel_features: 'identity' 5 | train: 6 | total_steps: 277800 7 | train_batch_size: 256 8 | learning_rate: 9 | type: 'step' 10 | init_learning_rate: 0.28 11 | learning_rate_levels: [0.028, 0.0028] 12 | learning_rate_steps: [263910, 273170] 13 | l2_weight_decay: 0.00003 14 | batch_norm_activation: 15 | use_sync_bn: true 16 | activation: 'swish' 17 | retinanet_head: 18 | num_filters: 48 19 | use_separable_conv: true 20 | spinenet_mbconv: 21 | model_id: '49' 22 | anchor: 23 | anchor_size: 3.0 24 | retinanet_parser: 25 | output_size: [384, 384] 26 | aug_scale_min: 0.5 27 | aug_scale_max: 2.0 28 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet143_classification.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-143 ImageNet classification: 2 | # - 79.0 Top-1, 94.4 Top-5 accuracy with init_drop_connect_rate null, activation 'relu', label_smoothing 0.0 3 | # - 80.1 Top-1, 95.0 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | parser: 'classification_parser' 8 | num_classes: 1001 9 | train: 10 | total_steps: 62557 11 | train_batch_size: 4096 12 | learning_rate: 13 | type: cosine 14 | warmup_steps: 1564 15 | init_learning_rate: 1.6 16 | l2_weight_decay: 0.0001 17 | label_smoothing: 0.1 18 | batch_norm_activation: 19 | batch_norm_epsilon: 1.0e-05 20 | batch_norm_momentum: 0.9 21 | use_sync_bn: false 22 | activation: 'swish' 23 | spinenet: 24 | model_id: '143' 25 | init_drop_connect_rate: 0.2 26 | classification_head: 27 | aggregation: 'all' 28 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet143_mrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-143 + Mask R-CNN: box mAP: 48.79, mask mAP: 42.71 2 | architecture: 3 | backbone: 'spinenet' 4 | min_level: 3 5 | max_level: 7 6 | multilevel_features: 'identity' 7 | parser: 'maskrcnn_parser' 8 | train: 9 | total_steps: 162050 10 | train_batch_size: 256 11 | learning_rate: 12 | warmup_steps: 2000 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [148160, 157420] 16 | l2_weight_decay: 0.00004 17 | anchor: 18 | anchor_size: 3.0 19 | batch_norm_activation: 20 | batch_norm_epsilon: 0.001 21 | batch_norm_momentum: 0.99 22 | use_sync_bn: true 23 | spinenet: 24 | model_id: '143' 25 | maskrcnn_parser: 26 | output_size: [1280, 1280] 27 | aug_scale_min: 0.5 28 | aug_scale_max: 2.0 29 | rpn_head: 30 | use_batch_norm: true 31 | frcnn_head: 32 | num_convs: 4 33 | num_fcs: 1 34 | use_batch_norm: true 35 | mrcnn_head: 36 | use_batch_norm: true 37 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet143_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-143 + RetinaNet: 2 | # - 48.0% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420], aug_scale_min 0.5, aug_scale_max 2.0 3 | # - 50.6% mAP with init_drop_connect_rate: 0.16, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870], aug_scale_min 0.1, aug_scale_max 1.9 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | train: 8 | total_steps: 231500 9 | train_batch_size: 256 10 | learning_rate: 11 | warmup_steps: 2000 12 | init_learning_rate: 0.28 13 | learning_rate_levels: [0.028, 0.0028] 14 | learning_rate_steps: [217610, 226870] 15 | l2_weight_decay: 0.00004 16 | gradient_clip_norm: 10.0 17 | batch_norm_activation: 18 | use_sync_bn: true 19 | activation: 'swish' 20 | spinenet: 21 | model_id: '143' 22 | init_drop_connect_rate: 0.16 23 | retinanet_parser: 24 | output_size: [1280, 1280] 25 | aug_scale_min: 0.1 26 | aug_scale_max: 1.9 27 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet190_classification.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-190 ImageNet classification: 2 | # - 80.8 Top-1, 95.3 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 3 | architecture: 4 | backbone: 'spinenet' 5 | multilevel_features: 'identity' 6 | parser: 'classification_parser' 7 | num_classes: 1001 8 | train: 9 | total_steps: 62557 10 | train_batch_size: 4096 11 | learning_rate: 12 | type: cosine 13 | warmup_steps: 1564 14 | init_learning_rate: 1.6 15 | l2_weight_decay: 0.0001 16 | label_smoothing: 0.1 17 | batch_norm_activation: 18 | batch_norm_epsilon: 1.0e-05 19 | batch_norm_momentum: 0.9 20 | use_sync_bn: false 21 | activation: 'swish' 22 | spinenet: 23 | model_id: '190' 24 | init_drop_connect_rate: 0.2 25 | classification_head: 26 | aggregation: 'all' 27 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet190_mrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-190 + Mask R-CNN with stochastic depth and swish activation. 52.0 mAP. 2 | architecture: 3 | backbone: 'spinenet' 4 | min_level: 3 5 | max_level: 7 6 | multilevel_features: 'identity' 7 | parser: 'maskrcnn_parser' 8 | train: 9 | total_steps: 187600 10 | train_batch_size: 256 11 | learning_rate: 12 | warmup_steps: 2000 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [173530, 182910] 16 | l2_weight_decay: 0.00004 17 | anchor: 18 | anchor_size: 5.0 19 | batch_norm_activation: 20 | batch_norm_epsilon: 0.001 21 | batch_norm_momentum: 0.99 22 | use_sync_bn: true 23 | activation: 'swish' 24 | spinenet: 25 | model_id: '190' 26 | init_drop_connect_rate: 0.2 27 | maskrcnn_parser: 28 | output_size: [1536, 1536] 29 | aug_scale_min: 0.1 30 | aug_scale_max: 2.0 31 | rpn_head: 32 | num_convs: 7 33 | num_filters: 384 34 | use_batch_norm: true 35 | frcnn_head: 36 | num_convs: 7 37 | num_filters: 384 38 | num_fcs: 1 39 | use_batch_norm: true 40 | mrcnn_head: 41 | num_convs: 7 42 | num_filters: 384 43 | use_batch_norm: true 44 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet190_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-190 + RetinaNet: 2 | # - 52.0% mAP with init_drop_connect_rate: 0.2, activation: 'swish' 3 | architecture: 4 | backbone: 'spinenet' 5 | multilevel_features: 'identity' 6 | train: 7 | total_steps: 185200 8 | train_batch_size: 256 9 | learning_rate: 10 | warmup_steps: 2000 11 | init_learning_rate: 0.3 12 | learning_rate_levels: [0.03, 0.003] 13 | learning_rate_steps: [171310, 180570] 14 | l2_weight_decay: 0.00004 15 | gradient_clip_norm: 10.0 16 | batch_norm_activation: 17 | use_sync_bn: true 18 | activation: 'swish' 19 | spinenet: 20 | model_id: '190' 21 | init_drop_connect_rate: 0.2 22 | retinanet_head: 23 | num_filters: 512 24 | num_convs: 7 25 | retinanet_parser: 26 | output_size: [1280, 1280] 27 | aug_scale_min: 0.1 28 | aug_scale_max: 1.9 29 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet49S_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-49S + RetinaNet: 2 | # - 39.7% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420] 3 | # - 41.5% mAP with init_drop_connect_rate: 0.2, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870] 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | train: 8 | total_steps: 231500 9 | train_batch_size: 256 10 | learning_rate: 11 | warmup_steps: 2000 12 | init_learning_rate: 0.28 13 | learning_rate_levels: [0.028, 0.0028] 14 | learning_rate_steps: [217610, 226870] 15 | l2_weight_decay: 0.00004 16 | gradient_clip_norm: 10.0 17 | batch_norm_activation: 18 | use_sync_bn: true 19 | activation: 'swish' 20 | spinenet: 21 | model_id: '49S' 22 | init_drop_connect_rate: 0.2 23 | retinanet_head: 24 | num_filters: 128 25 | anchor: 26 | anchor_size: 3.0 27 | retinanet_parser: 28 | output_size: [640, 640] 29 | aug_scale_min: 0.5 30 | aug_scale_max: 2.0 31 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet49_classification.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-49 ImageNet classification: 2 | # - 77.0 Top-1, 93.3 Top-5 accuracy with init_drop_connect_rate null, activation 'relu', label_smoothing 0.0 3 | # - 78.1 Top-1, 94.0 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | parser: 'classification_parser' 8 | num_classes: 1001 9 | train: 10 | total_steps: 62557 11 | train_batch_size: 4096 12 | learning_rate: 13 | type: cosine 14 | warmup_steps: 1564 15 | init_learning_rate: 1.6 16 | l2_weight_decay: 0.0001 17 | label_smoothing: 0.1 18 | batch_norm_activation: 19 | batch_norm_epsilon: 1.0e-05 20 | batch_norm_momentum: 0.9 21 | use_sync_bn: false 22 | activation: 'swish' 23 | spinenet: 24 | model_id: '49' 25 | init_drop_connect_rate: 0.2 26 | classification_head: 27 | aggregation: 'all' 28 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet49_mrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-49 + Mask R-CNN: box mAP: 42.93, mask mAP: 38.07 2 | architecture: 3 | backbone: 'spinenet' 4 | min_level: 3 5 | max_level: 7 6 | multilevel_features: 'identity' 7 | parser: 'maskrcnn_parser' 8 | train: 9 | total_steps: 162050 10 | train_batch_size: 256 11 | learning_rate: 12 | warmup_steps: 2000 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [148160, 157420] 16 | l2_weight_decay: 0.00004 17 | anchor: 18 | anchor_size: 3.0 19 | batch_norm_activation: 20 | batch_norm_epsilon: 0.001 21 | batch_norm_momentum: 0.99 22 | use_sync_bn: true 23 | spinenet: 24 | model_id: '49' 25 | maskrcnn_parser: 26 | output_size: [640, 640] 27 | aug_scale_min: 0.5 28 | aug_scale_max: 2.0 29 | rpn_head: 30 | use_batch_norm: true 31 | frcnn_head: 32 | num_convs: 4 33 | num_fcs: 1 34 | use_batch_norm: true 35 | mrcnn_head: 36 | use_batch_norm: true 37 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet49_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-49 + RetinaNet: 2 | # - 42.8% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420] 3 | # - 44.3% mAP with init_drop_connect_rate: 0.2, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870] 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | train: 8 | total_steps: 231500 9 | train_batch_size: 256 10 | learning_rate: 11 | warmup_steps: 2000 12 | init_learning_rate: 0.28 13 | learning_rate_levels: [0.028, 0.0028] 14 | learning_rate_steps: [217610, 226870] 15 | l2_weight_decay: 0.00004 16 | gradient_clip_norm: 10.0 17 | batch_norm_activation: 18 | use_sync_bn: true 19 | activation: 'swish' 20 | spinenet: 21 | model_id: '49' 22 | init_drop_connect_rate: 0.2 23 | anchor: 24 | anchor_size: 3.0 25 | retinanet_parser: 26 | output_size: [640, 640] 27 | aug_scale_min: 0.5 28 | aug_scale_max: 2.0 29 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet49s_mrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-49S + Mask R-CNN: box mAP: 39.31, mask mAP: 34.83 2 | architecture: 3 | backbone: 'spinenet' 4 | min_level: 3 5 | max_level: 7 6 | multilevel_features: 'identity' 7 | parser: 'maskrcnn_parser' 8 | train: 9 | total_steps: 162050 10 | train_batch_size: 256 11 | learning_rate: 12 | warmup_steps: 2000 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [148160, 157420] 16 | l2_weight_decay: 0.00004 17 | anchor: 18 | anchor_size: 3.0 19 | batch_norm_activation: 20 | batch_norm_epsilon: 0.001 21 | batch_norm_momentum: 0.99 22 | use_sync_bn: true 23 | spinenet: 24 | model_id: '49S' 25 | maskrcnn_parser: 26 | output_size: [640, 640] 27 | aug_scale_min: 0.5 28 | aug_scale_max: 2.0 29 | rpn_head: 30 | num_filters: 128 31 | use_batch_norm: true 32 | frcnn_head: 33 | num_convs: 4 34 | num_filters: 128 35 | num_fcs: 1 36 | fc_dims: 512 37 | use_batch_norm: true 38 | mrcnn_head: 39 | num_filters: 128 40 | use_batch_norm: true 41 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet96_classification.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-96 ImageNet classification: 2 | # - 78.2 Top-1, 94.0 Top-5 accuracy with init_drop_connect_rate null, activation 'relu', label_smoothing 0.0 3 | # - 79.4 Top-1, 94.6 Top-5 accuracy with init_drop_connect_rate 0.2, activation 'swish', label_smoothing 0.1 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | parser: 'classification_parser' 8 | num_classes: 1001 9 | train: 10 | total_steps: 62557 11 | train_batch_size: 4096 12 | learning_rate: 13 | type: cosine 14 | warmup_steps: 1564 15 | init_learning_rate: 1.6 16 | l2_weight_decay: 0.0001 17 | label_smoothing: 0.1 18 | batch_norm_activation: 19 | batch_norm_epsilon: 1.0e-05 20 | batch_norm_momentum: 0.9 21 | use_sync_bn: false 22 | activation: 'swish' 23 | spinenet: 24 | model_id: '96' 25 | init_drop_connect_rate: 0.2 26 | classification_head: 27 | aggregation: 'all' 28 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet96_mrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-96 + Mask R-CNN: box mAP: 47.18, mask mAP: 41.53 2 | architecture: 3 | backbone: 'spinenet' 4 | min_level: 3 5 | max_level: 7 6 | multilevel_features: 'identity' 7 | parser: 'maskrcnn_parser' 8 | train: 9 | total_steps: 162050 10 | train_batch_size: 256 11 | learning_rate: 12 | warmup_steps: 2000 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [148160, 157420] 16 | l2_weight_decay: 0.00004 17 | anchor: 18 | anchor_size: 3.0 19 | batch_norm_activation: 20 | batch_norm_epsilon: 0.001 21 | batch_norm_momentum: 0.99 22 | use_sync_bn: true 23 | spinenet: 24 | model_id: '96' 25 | maskrcnn_parser: 26 | output_size: [1024, 1024] 27 | aug_scale_min: 0.5 28 | aug_scale_max: 2.0 29 | rpn_head: 30 | use_batch_norm: true 31 | frcnn_head: 32 | num_convs: 4 33 | num_fcs: 1 34 | use_batch_norm: true 35 | mrcnn_head: 36 | use_batch_norm: true 37 | -------------------------------------------------------------------------------- /models/official/detection/configs/spinenet/spinenet96_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-96 + RetinaNet: 2 | # - 46.7% mAP with init_drop_connect_rate null, activation 'relu', total_steps 162050, learning_rate_steps [148160, 157420] 3 | # - 48.5% mAP with init_drop_connect_rate: 0.2, activation: 'swish', total_steps 231500, learning_rate_steps [217610, 226870] 4 | 5 | architecture: 6 | backbone: 'spinenet' 7 | multilevel_features: 'identity' 8 | train: 9 | total_steps: 231500 10 | train_batch_size: 256 11 | learning_rate: 12 | warmup_steps: 2000 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [217610, 226870] 16 | l2_weight_decay: 0.00004 17 | gradient_clip_norm: 10.0 18 | batch_norm_activation: 19 | use_sync_bn: true 20 | activation: 'swish' 21 | spinenet: 22 | model_id: '96' 23 | init_drop_connect_rate: 0.2 24 | anchor: 25 | anchor_size: 3.0 26 | retinanet_parser: 27 | output_size: [1024, 1024] 28 | aug_scale_min: 0.5 29 | aug_scale_max: 2.0 30 | -------------------------------------------------------------------------------- /models/official/detection/configs/yaml/retinanet_autoaugment.yaml: -------------------------------------------------------------------------------- 1 | # ---------- TRAINING PARAMETERS ---------- 2 | # AutoAugment achieves best results when trained for long and when not using a pretrained 3 | # checkpoint for the backbone. 4 | # Right now the code uses a pretrained checkpoint, but a future version will change this. 5 | # To train the ResNet 101 or ResNet 200 version, simply change the 'resnet_depth' to 101 or 200. 6 | # Expected accuracy on ResNet 50 with using autoaugment: 38.0 7 | # Expected accuracy on ResNet 50 without using autoaugment: 36.4 8 | train: 9 | total_steps: 277200 10 | learning_rate: 11 | init_learning_rate: 0.08 12 | learning_rate_levels: [0.008, 0.0008] 13 | learning_rate_steps: [220000, 258700] 14 | 15 | resnet: 16 | resnet_depth: 50 17 | 18 | retinanet_parser: 19 | aug_policy: 'v0' 20 | aug_scale_min: 0.8 21 | aug_scale_max: 1.2 22 | -------------------------------------------------------------------------------- /models/official/detection/configs/yaml/retinanet_nasfpn.yaml: -------------------------------------------------------------------------------- 1 | # ---------- RetianNet + NAS-FPN ---------- 2 | # Expected accuracy with using NAS-FPN l3-l7 and image size 640x640: 39.5 3 | train: 4 | total_steps: 90000 5 | learning_rate: 6 | init_learning_rate: 0.08 7 | learning_rate_levels: [0.008, 0.0008] 8 | learning_rate_steps: [60000, 80000] 9 | 10 | architecture: 11 | multilevel_features: 'nasfpn' 12 | 13 | nasfpn: 14 | fpn_feat_dims: 256 15 | min_level: 3 16 | max_level: 7 17 | num_repeats: 5 18 | use_separable_conv: False 19 | 20 | retinanet_parser: 21 | aug_scale_min: 0.8 22 | aug_scale_max: 1.2 23 | -------------------------------------------------------------------------------- /models/official/detection/dataloader/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/dataloader/mode_keys.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Standard names for input dataloader modes. 16 | 17 | The following standard keys are defined: 18 | 19 | * `TRAIN`: training mode. 20 | * `EVAL`: evaluation mode. 21 | * `PREDICT`: prediction mode. 22 | * `PREDICT_WITH_GT`: prediction mode with groundtruths in returned variables. 23 | """ 24 | 25 | from __future__ import absolute_import 26 | from __future__ import division 27 | from __future__ import print_function 28 | 29 | 30 | TRAIN = 'train' 31 | EVAL = 'eval' 32 | PREDICT = 'predict' 33 | PREDICT_WITH_GT = 'predict_with_gt' 34 | -------------------------------------------------------------------------------- /models/official/detection/datasets/coco_label_map.csv: -------------------------------------------------------------------------------- 1 | 1:person 2 | 2:bicycle 3 | 3:car 4 | 4:motorcycle 5 | 5:airplane 6 | 6:bus 7 | 7:train 8 | 8:truck 9 | 9:boat 10 | 10:traffic light 11 | 11:fire hydrant 12 | 13:stop sign 13 | 14:parking meter 14 | 15:bench 15 | 16:bird 16 | 17:cat 17 | 18:dog 18 | 19:horse 19 | 20:sheep 20 | 21:cow 21 | 22:elephant 22 | 23:bear 23 | 24:zebra 24 | 25:giraffe 25 | 27:backpack 26 | 28:umbrella 27 | 31:handbag 28 | 32:tie 29 | 33:suitcase 30 | 34:frisbee 31 | 35:skis 32 | 36:snowboard 33 | 37:sports ball 34 | 38:kite 35 | 39:baseball bat 36 | 40:baseball glove 37 | 41:skateboard 38 | 42:surfboard 39 | 43:tennis racket 40 | 44:bottle 41 | 46:wine glass 42 | 47:cup 43 | 48:fork 44 | 49:knife 45 | 50:spoon 46 | 51:bowl 47 | 52:banana 48 | 53:apple 49 | 54:sandwich 50 | 55:orange 51 | 56:broccoli 52 | 57:carrot 53 | 58:hot dog 54 | 59:pizza 55 | 60:donut 56 | 61:cake 57 | 62:chair 58 | 63:couch 59 | 64:potted plant 60 | 65:bed 61 | 67:dining table 62 | 70:toilet 63 | 72:tv 64 | 73:laptop 65 | 74:mouse 66 | 75:remote 67 | 76:keyboard 68 | 77:cell phone 69 | 78:microwave 70 | 79:oven 71 | 80:toaster 72 | 81:sink 73 | 82:refrigerator 74 | 84:book 75 | 85:clock 76 | 86:vase 77 | 87:scissors 78 | 88:teddy bear 79 | 89:hair drier 80 | 90:toothbrush 81 | -------------------------------------------------------------------------------- /models/official/detection/evaluation/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/executor/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/modeling/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/modeling/architecture/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/modeling/architecture/identity.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Identity Fn that forwards the input features.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | 22 | class Identity(object): 23 | """Identity function that forwards the input features.""" 24 | 25 | def __call__(self, features, is_training=False): 26 | """Only forwards the input features.""" 27 | return features 28 | 29 | -------------------------------------------------------------------------------- /models/official/detection/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/configs/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/configs/factory.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Factory to provide Attribute-Mask R-CNN model configs.""" 16 | 17 | from projects.fashionpedia.configs import model_config 18 | from hyperparameters import params_dict 19 | 20 | 21 | def config_generator(model): 22 | """Model function generator.""" 23 | if model == 'attribute_mask_rcnn': 24 | default_config = model_config.CFG 25 | restrictions = model_config.RESTRICTIONS 26 | else: 27 | raise ValueError('Model %s is not supported.' % model) 28 | 29 | return params_dict.ParamsDict(default_config, restrictions) 30 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/configs/yaml/r101fpn_amrcnn.yaml: -------------------------------------------------------------------------------- 1 | # ResNet-101 FPN + Attribute-Mask R-CNN 2 | # 3x schedule (~95 Fashionpedia epochs) 3 | # Box AP (IoU / IoU+F1): 44.9 / 32.8 4 | # Mask AP (IoU / IoU+F1): 40.7 / 31.4 5 | architecture: 6 | backbone: 'resnet' 7 | multilevel_features: 'fpn' 8 | train: 9 | total_steps: 16875 10 | train_batch_size: 256 11 | learning_rate: 12 | type: 'step' 13 | init_learning_rate: 0.32 14 | learning_rate_levels: [0.032, 0.0032] 15 | learning_rate_steps: [15000, 16250] 16 | batch_norm_activation: 17 | use_sync_bn: true 18 | resnet: 19 | resnet_depth: 101 20 | attribute_maskrcnn_parser: 21 | output_size: [1024, 1024] 22 | aug_scale_min: 0.5 23 | aug_scale_max: 2.0 24 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/configs/yaml/r50fpn_amrcnn.yaml: -------------------------------------------------------------------------------- 1 | # ResNet-50 FPN + Attribute-Mask R-CNN 2 | # 3x schedule (~95 Fashionpedia epochs) 3 | # Box AP (IoU / IoU+F1): 43.4 / 30.7 4 | # Mask AP (IoU / IoU+F1): 39.2 / 29.5 5 | architecture: 6 | backbone: 'resnet' 7 | multilevel_features: 'fpn' 8 | train: 9 | total_steps: 16875 10 | train_batch_size: 256 11 | learning_rate: 12 | type: 'step' 13 | init_learning_rate: 0.32 14 | learning_rate_levels: [0.032, 0.0032] 15 | learning_rate_steps: [15000, 16250] 16 | batch_norm_activation: 17 | use_sync_bn: true 18 | resnet: 19 | resnet_depth: 50 20 | attribute_maskrcnn_parser: 21 | output_size: [1024, 1024] 22 | aug_scale_min: 0.5 23 | aug_scale_max: 2.0 24 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/configs/yaml/spinenet143_amrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-143 + Attribute-Mask R-CNN 2 | # 6x schedule (~189 Fashionpedia epochs) 3 | # Box AP (IoU / IoU+F1): 48.7 / 35.7 4 | # Mask AP (IoU / IoU+F1): 43.1 / 33.3 5 | architecture: 6 | backbone: 'spinenet' 7 | multilevel_features: 'identity' 8 | train: 9 | total_steps: 33750 10 | train_batch_size: 256 11 | learning_rate: 12 | type: 'step' 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [31875, 33125] 16 | batch_norm_activation: 17 | use_sync_bn: true 18 | activation: 'swish' 19 | spinenet: 20 | model_id: '143' 21 | init_drop_connect_rate: 0.2 22 | attribute_maskrcnn_parser: 23 | output_size: [1280, 1280] 24 | aug_scale_min: 0.5 25 | aug_scale_max: 2.0 26 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/configs/yaml/spinenet49_amrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-49 + Attribute-Mask R-CNN 2 | # 6x schedule (~189 Fashionpedia epochs) 3 | # Box AP (IoU / IoU+F1): 43.7 / 32.4 4 | # Mask AP (IoU / IoU+F1): 39.6 / 31.4 5 | architecture: 6 | backbone: 'spinenet' 7 | multilevel_features: 'identity' 8 | train: 9 | total_steps: 33750 10 | train_batch_size: 256 11 | learning_rate: 12 | type: 'step' 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [31875, 33125] 16 | batch_norm_activation: 17 | use_sync_bn: true 18 | activation: 'swish' 19 | spinenet: 20 | model_id: '49' 21 | init_drop_connect_rate: 0.2 22 | attribute_maskrcnn_parser: 23 | output_size: [1024, 1024] 24 | aug_scale_min: 0.5 25 | aug_scale_max: 2.0 26 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/configs/yaml/spinenet96_amrcnn.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-96 + Attribute-Mask R-CNN 2 | # 6x schedule (~189 Fashionpedia epochs) 3 | # Box AP (IoU / IoU+F1): 46.4 / 34.0 4 | # Mask AP (IoU / IoU+F1): 41.2 / 31.8 5 | architecture: 6 | backbone: 'spinenet' 7 | multilevel_features: 'identity' 8 | train: 9 | total_steps: 33750 10 | train_batch_size: 256 11 | learning_rate: 12 | type: 'step' 13 | init_learning_rate: 0.28 14 | learning_rate_levels: [0.028, 0.0028] 15 | learning_rate_steps: [31875, 33125] 16 | batch_norm_activation: 17 | use_sync_bn: true 18 | activation: 'swish' 19 | spinenet: 20 | model_id: '96' 21 | init_drop_connect_rate: 0.2 22 | attribute_maskrcnn_parser: 23 | output_size: [1024, 1024] 24 | aug_scale_min: 0.5 25 | aug_scale_max: 2.0 26 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/dataloader/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/dataset/fashionpedia_label_map.csv: -------------------------------------------------------------------------------- 1 | 1:shirt, blouse 2 | 2:top, t-shirt, sweatshirt 3 | 3:sweater 4 | 4:cardigan 5 | 5:jacket 6 | 6:vest 7 | 7:pants 8 | 8:shorts 9 | 9:skirt 10 | 10:coat 11 | 11:dress 12 | 12:jumpsuit 13 | 13:cape 14 | 14:glasses 15 | 15:hat 16 | 16:headband, head covering, hair accessory 17 | 17:tie 18 | 18:glove 19 | 19:watch 20 | 20:belt 21 | 21:leg warmer 22 | 22:tights, stockings 23 | 23:sock 24 | 24:shoe 25 | 25:bag, wallet 26 | 26:scarf 27 | 27:umbrella 28 | 28:hood 29 | 29:collar 30 | 30:lapel 31 | 31:epaulette 32 | 32:sleeve 33 | 33:pocket 34 | 34:neckline 35 | 35:buckle 36 | 36:zipper 37 | 37:applique 38 | 38:bead 39 | 39:bow 40 | 40:flower 41 | 41:fringe 42 | 42:ribbon 43 | 43:rivet 44 | 44:ruffle 45 | 45:sequin 46 | 46:tassel 47 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/modeling/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/modeling/architecture/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/modeling/factory.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Factory to build Attribute-Mask R-CNN model.""" 16 | 17 | from projects.fashionpedia.modeling import model 18 | 19 | 20 | def model_generator(params): 21 | """Model function generator.""" 22 | if params.type == 'attribute_mask_rcnn': 23 | model_fn = model.AttributeMaskrcnnModel(params) 24 | else: 25 | raise ValueError('Model %s is not supported.'% params.type) 26 | 27 | return model_fn 28 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/serving/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/fashionpedia/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/self_training/configs/coco_spinenet143_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-143 + RetinaNet: 2 | # Config is used to run eval on the SpineNet-143 checkpoint trained with semi-supervised learning 3 | # Performance: 52.2 AP without SoftNMS, 52.4 with SoftNMS 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | train: 8 | total_steps: 231500 9 | train_batch_size: 256 10 | learning_rate: 11 | warmup_steps: 2000 12 | init_learning_rate: 0.28 13 | learning_rate_levels: [0.028, 0.0028] 14 | learning_rate_steps: [217610, 226870] 15 | l2_weight_decay: 0.00004 16 | gradient_clip_norm: 10.0 17 | batch_norm_activation: 18 | use_sync_bn: true 19 | activation: 'swish' 20 | batch_norm_epsilon: 0.001 21 | spinenet: 22 | model_id: '143' 23 | init_drop_connect_rate: 0.16 24 | retinanet_parser: 25 | output_size: [1280, 1280] 26 | aug_scale_min: 0.1 27 | aug_scale_max: 1.9 28 | -------------------------------------------------------------------------------- /models/official/detection/projects/self_training/configs/coco_spinenet190_retinanet.yaml: -------------------------------------------------------------------------------- 1 | # SpineNet-190 + RetinaNet: 2 | # Config is used to run eval on the SpineNet-190 checkpoint trained with semi-supervised learning 3 | # Performance: 53.9 AP without SoftNMS, 54.2 with SoftNMS 4 | architecture: 5 | backbone: 'spinenet' 6 | multilevel_features: 'identity' 7 | train: 8 | total_steps: 185200 9 | train_batch_size: 256 10 | learning_rate: 11 | warmup_steps: 2000 12 | init_learning_rate: 0.3 13 | learning_rate_levels: [0.03, 0.003] 14 | learning_rate_steps: [171310, 180570] 15 | l2_weight_decay: 0.00004 16 | gradient_clip_norm: 10.0 17 | batch_norm_activation: 18 | use_sync_bn: true 19 | activation: 'swish' 20 | batch_norm_epsilon: 0.001 21 | spinenet: 22 | model_id: '190' 23 | init_drop_connect_rate: 0.2 24 | retinanet_head: 25 | num_filters: 512 26 | num_convs: 7 27 | retinanet_parser: 28 | output_size: [1280, 1280] 29 | aug_scale_min: 0.1 30 | aug_scale_max: 1.9 31 | -------------------------------------------------------------------------------- /models/official/detection/projects/vild/preprocessing/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/projects/vild/requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow 2 | pycocotools 3 | lvis 4 | opencv-python -------------------------------------------------------------------------------- /models/official/detection/serving/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/detection/utils/object_detection/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/efficientnet/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/efficientnet/condconv/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/efficientnet/edgetpu/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/efficientnet/g3doc/condconv-layer.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/efficientnet/g3doc/condconv-layer.png -------------------------------------------------------------------------------- /models/official/efficientnet/g3doc/efficientnet-edgetpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/efficientnet/g3doc/efficientnet-edgetpu.png -------------------------------------------------------------------------------- /models/official/efficientnet/g3doc/flops.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/efficientnet/g3doc/flops.png -------------------------------------------------------------------------------- /models/official/efficientnet/g3doc/lite-float-gpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/efficientnet/g3doc/lite-float-gpu.png -------------------------------------------------------------------------------- /models/official/efficientnet/g3doc/lite-quant-cpu.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/efficientnet/g3doc/lite-quant-cpu.png -------------------------------------------------------------------------------- /models/official/efficientnet/g3doc/lite-quant-size.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/efficientnet/g3doc/lite-quant-size.png -------------------------------------------------------------------------------- /models/official/efficientnet/g3doc/params.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/efficientnet/g3doc/params.png -------------------------------------------------------------------------------- /models/official/efficientnet/lite/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/efficientnet/tpu/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/README.md: -------------------------------------------------------------------------------- 1 | # Try it 2 | 3 | Try to run our pre-trained COCO Mask R-CNN using [Colab](https://colab.sandbox.google.com/github/tensorflow/tpu/blob/master/models/official/mask_rcnn/mask_rcnn_demo.ipynb). 4 | 5 | # Installing extra packages 6 | 7 | Mask R-CNN requires a few extra packages. We can install them now: 8 | 9 | ``` 10 | sudo apt-get install -y python-tk && \ 11 | pip3 install --user Cython matplotlib opencv-python-headless pyyaml Pillow && \ 12 | pip3 install --user 'git+https://github.com/cocodataset/cocoapi#egg=pycocotools&subdirectory=PythonAPI' 13 | ``` 14 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/configs/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/configs/cloud/v2-128.yaml: -------------------------------------------------------------------------------- 1 | # ---------- MODEL PARAMETERS ------------- 2 | backbone: 'resnet50' 3 | num_cores: 128 4 | # ---------- TRAINING PARAMETERS ---------- 5 | train_batch_size: 512 6 | init_learning_rate: 0.24 7 | warmup_learning_rate: 0.0067 8 | warmup_steps: 1600 9 | learning_rate_levels: [0.024, 0.0024, 0.00024] 10 | learning_rate_steps: [6000, 8000, 10000] 11 | total_steps: 11250 12 | global_gradient_clip_ratio: 0.02 13 | num_batch_norm_group: 1 14 | momentum: 0.95 15 | precision: 'bfloat16' 16 | # ---------- EVAL PARAMETERS -------------- 17 | eval_batch_size: 8 18 | eval_samples: 5000 19 | num_steps_per_eval: 2500 20 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/configs/cloud/v2-32.yaml: -------------------------------------------------------------------------------- 1 | # ---------- MODEL PARAMETERS ------------- 2 | backbone: 'resnet50' 3 | num_cores: 32 4 | # ---------- TRAINING PARAMETERS ---------- 5 | train_batch_size: 128 6 | init_learning_rate: 0.16 7 | warmup_learning_rate: 0.0067 8 | warmup_steps: 1000 9 | learning_rate_levels: [0.016, 0.0016] 10 | learning_rate_steps: [7500, 10000] 11 | total_steps: 11250 12 | global_gradient_clip_ratio: 0.02 13 | num_batch_norm_group: 1 14 | precision: 'bfloat16' 15 | # ---------- EVAL PARAMETERS -------------- 16 | eval_batch_size: 8 17 | eval_samples: 5000 18 | num_steps_per_eval: 2500 19 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/configs/cloud/v2-8.yaml: -------------------------------------------------------------------------------- 1 | # ---------- MODEL PARAMETERS ------------- 2 | backbone: 'resnet50' 3 | num_cores: 8 4 | # ---------- TRAINING PARAMETERS ---------- 5 | train_batch_size: 32 6 | init_learning_rate: 0.04 7 | warmup_learning_rate: 0.0067 8 | warmup_steps: 500 9 | learning_rate_levels: [0.004, 0.0004] 10 | learning_rate_steps: [30000, 40000] 11 | total_steps: 45000 12 | precision: 'bfloat16' 13 | # ---------- EVAL PARAMETERS -------------- 14 | eval_batch_size: 8 15 | eval_samples: 5000 16 | num_steps_per_eval: 2500 17 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/configs/cloud/v3-128.yaml: -------------------------------------------------------------------------------- 1 | # ---------- MODEL PARAMETERS ------------- 2 | backbone: 'resnet50' 3 | num_cores: 128 4 | # ---------- TRAINING PARAMETERS ---------- 5 | train_batch_size: 512 6 | init_learning_rate: 0.24 7 | warmup_learning_rate: 0.0067 8 | warmup_steps: 1600 9 | learning_rate_levels: [0.024, 0.0024, 0.00024] 10 | learning_rate_steps: [6000, 8000, 10000] 11 | total_steps: 11250 12 | global_gradient_clip_ratio: 0.02 13 | num_batch_norm_group: 1 14 | momentum: 0.95 15 | precision: 'bfloat16' 16 | # ---------- EVAL PARAMETERS -------------- 17 | eval_batch_size: 8 18 | eval_samples: 5000 19 | num_steps_per_eval: 2500 20 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/configs/cloud/v3-32.yaml: -------------------------------------------------------------------------------- 1 | # ---------- MODEL PARAMETERS ------------- 2 | backbone: 'resnet50' 3 | num_cores: 32 4 | # ---------- TRAINING PARAMETERS ---------- 5 | train_batch_size: 128 6 | init_learning_rate: 0.16 7 | warmup_learning_rate: 0.0067 8 | warmup_steps: 1000 9 | learning_rate_levels: [0.016, 0.0016] 10 | learning_rate_steps: [7500, 10000] 11 | total_steps: 11250 12 | global_gradient_clip_ratio: 0.02 13 | num_batch_norm_group: 1 14 | precision: 'bfloat16' 15 | # ---------- EVAL PARAMETERS -------------- 16 | eval_batch_size: 8 17 | eval_samples: 5000 18 | num_steps_per_eval: 2500 19 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/configs/cloud/v3-8.yaml: -------------------------------------------------------------------------------- 1 | # ---------- MODEL PARAMETERS ------------- 2 | backbone: 'resnet50' 3 | num_cores: 8 4 | # ---------- TRAINING PARAMETERS ---------- 5 | train_batch_size: 64 6 | init_learning_rate: 0.08 7 | warmup_learning_rate: 0.0067 8 | warmup_steps: 500 9 | learning_rate_levels: [0.008, 0.0008] 10 | learning_rate_steps: [15000, 20000] 11 | total_steps: 22500 12 | precision: 'bfloat16' 13 | # ---------- EVAL PARAMETERS -------------- 14 | eval_batch_size: 8 15 | eval_samples: 5000 16 | num_steps_per_eval: 2500 17 | -------------------------------------------------------------------------------- /models/official/mask_rcnn/object_detection/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/mnasnet/configs/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/mnasnet/configs/cloud/gpu.yaml: -------------------------------------------------------------------------------- 1 | use_tpu: False 2 | train_steps: 3503192 # 1281167 * 350 / train_batch_size 3 | train_batch_size: 128 # 1024 / 8 4 | eval_batch_size: 128 # 1024 / 8 5 | model_name: 'mnasnet-a1' 6 | dropout_rate: null 7 | depth_multiplier: null 8 | -------------------------------------------------------------------------------- /models/official/mnasnet/configs/cloud/v2-32.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 109474 2 | train_batch_size: 4096 3 | eval_batch_size: 256 4 | iterations_per_loop: 100 5 | skip_host_call: false 6 | model_name: 'mnasnet-a1' 7 | dropout_rate: null 8 | depth_multiplier: null 9 | use_keras: true 10 | precision: 'float32' 11 | -------------------------------------------------------------------------------- /models/official/mnasnet/configs/cloud/v2-8.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 437899 2 | train_batch_size: 1024 3 | eval_batch_size: 1024 4 | iterations_per_loop: 1251 5 | skip_host_call: True 6 | model_name: 'mnasnet-a1' 7 | dropout_rate: null 8 | depth_multiplier: null 9 | -------------------------------------------------------------------------------- /models/official/mnasnet/configs/cloud/v3-32.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 109474 2 | train_batch_size: 4096 3 | eval_batch_size: 256 4 | iterations_per_loop: 100 5 | skip_host_call: false 6 | model_name: 'mnasnet-a1' 7 | dropout_rate: null 8 | depth_multiplier: null 9 | use_keras: true 10 | precision: 'float32' 11 | -------------------------------------------------------------------------------- /models/official/mnasnet/configs/cloud/v3-8.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 437899 2 | train_batch_size: 1024 3 | eval_batch_size: 1024 4 | iterations_per_loop: 1251 5 | skip_host_call: True 6 | model_name: 'mnasnet-a1' 7 | dropout_rate: null 8 | depth_multiplier: null 9 | -------------------------------------------------------------------------------- /models/official/mnasnet/g3doc/mnasnet_vs_mobilenetv2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/mnasnet/g3doc/mnasnet_vs_mobilenetv2.png -------------------------------------------------------------------------------- /models/official/mnasnet/g3doc/mnasnet_vs_mobilenetv2_2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/mnasnet/g3doc/mnasnet_vs_mobilenetv2_2.png -------------------------------------------------------------------------------- /models/official/mnasnet/mixnet/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/mnasnet/mixnet/g3doc/mixnet-flops.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/mnasnet/mixnet/g3doc/mixnet-flops.png -------------------------------------------------------------------------------- /models/official/mnist/README.md: -------------------------------------------------------------------------------- 1 | `mnist_tpu.py` can be used to train a simple model on the MNIST dataset using 2 | a Cloud TPU. 3 | 4 | See https://cloud.google.com/tpu/docs/quickstart for more details. 5 | -------------------------------------------------------------------------------- /models/official/mobilenet/README.md: -------------------------------------------------------------------------------- 1 | # Cloud TPU Port of the MobileNet v1 model 2 | 3 | This is a straightforward port of the [MobileNet v1 model](https://arxiv.org/pdf/1704.04861.pdf). The code was based on the original version from the [tensorflow/models](https://github.com/tensorflow/models/tree/master/research/slim/nets) repository. 4 | 5 | The only adjustments have been to add the required code to enable using the 6 | TPUEstimator interface, along with the data processing pipeline for ImageNet. 7 | 8 | ## Running the model 9 | 10 | Assuming you have a version of ImageNet converted to the tfrecord format located 11 | at `gs://my-cloud-bucket/data/imagenet/`, you can run this model with the 12 | following command: 13 | 14 | ``` 15 | python mobilenet.py\ 16 | --alsologtostderr\ 17 | --master=$TPU_WORKER\ 18 | --data_dir=gs://my-cloud-bucket/data/imagenet\ 19 | --model_dir=gs://my-cloud-bucket/models/mobilenet/v0\ 20 | --num_shards=8\ 21 | --batch_size=1024\ 22 | --use_tpu=1\ 23 | ``` 24 | 25 | Note that the mobilenet network requires a large number of epochs to converge 26 | completely. 27 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v2-128.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 16384 3 | train_steps: 500000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 128 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v2-256.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 32768 3 | train_steps: 250000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 256 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v2-32.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 4096 3 | train_steps: 2000000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 32 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v2-512.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 65536 3 | train_steps: 125000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 512 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v2-8.yaml: -------------------------------------------------------------------------------- 1 | train_batch_size: 1024 2 | train_steps: 8000000 3 | eval_batch_size: 1024 4 | iterations_per_loop: 100 5 | num_cores: 8 6 | train_steps_per_eval: 2000 7 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-1024.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 131072 3 | train_steps: 62500 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 1024 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-128.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 16384 3 | train_steps: 500000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 128 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-2048.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 262144 3 | train_steps: 31250 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 2048 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-256.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 32768 3 | train_steps: 250000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 256 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-32.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 4096 3 | train_steps: 2000000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 32 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-512.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 65536 3 | train_steps: 125000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 512 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-64.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_batch_size: 8192 3 | train_steps: 1000000 4 | eval_batch_size: 1024 5 | iterations_per_loop: 100 6 | num_cores: 64 7 | train_steps_per_eval: 2000 8 | -------------------------------------------------------------------------------- /models/official/mobilenet/configs/cloud/v3-8.yaml: -------------------------------------------------------------------------------- 1 | train_batch_size: 1024 2 | train_steps: 8000000 3 | eval_batch_size: 1024 4 | iterations_per_loop: 100 5 | num_cores: 8 6 | train_steps_per_eval: 2000 7 | -------------------------------------------------------------------------------- /models/official/resnet/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/resnet/__init__.py -------------------------------------------------------------------------------- /models/official/resnet/benchmark/README.md: -------------------------------------------------------------------------------- 1 | # ResNet-50 Benchmark on Cloud TPU pods 2 | 3 | Submission for [DAWNBench](https://dawn.cs.stanford.edu/benchmark/index.html). 4 | 5 | This subdirectory contains the code needed to replicate the DAWNBench results 6 | for ResNet-50 on a Cloud TPU pod. The model used here is identical to the model 7 | in the parent directory. The only difference is that `resnet_benchmark.py` will 8 | generate checkpoints at every epoch and evaluate in a separate job. 9 | 10 | ## Instructions for training on single Cloud TPU 11 | 12 | 1. Add the top-level `/models` folder to the Python path with the command 13 | 14 | ``` 15 | export PYTHONPATH="$PYTHONPATH:/path/to/models" 16 | ``` 17 | 18 | 1. Train the model (roughly 90 epochs, 1 checkpoint per epoch): 19 | ``` 20 | python resnet_benchmark.py \ 21 | --tpu=[TPU NAME] \ 22 | --mode=train \ 23 | --data_dir=[PATH TO DATA] \ 24 | --model_dir=[PATH TO MODEL] \ 25 | --train_batch_size=1024 \ 26 | --train_steps=112590 \ 27 | --iterations_per_loop=1251 28 | ``` 29 | 30 | 1. Evaluate the model (run after train completes): 31 | ``` 32 | python resnet_benchmark.py \ 33 | --tpu=[TPU NAME] \ 34 | --mode=eval \ 35 | --data_dir=[PATH TO DATA] \ 36 | --model_dir=[PATH TO MODEL] 37 | ``` 38 | 39 | ## Instructions for training on a half TPU Pod 40 | 41 | Not yet available due to TPU Pod availability in Cloud. 42 | 43 | -------------------------------------------------------------------------------- /models/official/resnet/benchmark/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/models/official/resnet/benchmark/__init__.py -------------------------------------------------------------------------------- /models/official/resnet/configs/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/randaugment-32.yaml: -------------------------------------------------------------------------------- 1 | resnet_depth: 50 2 | train_steps: 56304 3 | train_batch_size: 4096 4 | eval_batch_size: 1024 5 | iterations_per_loop: 1000 6 | num_cores: 32 7 | skip_host_call: True 8 | augment_name: 'randaugment' 9 | randaug_num_layers: 2 10 | randaug_magnitude: 9 11 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/randaugment-8.yaml: -------------------------------------------------------------------------------- 1 | resnet_depth: 50 2 | train_steps: 225216 3 | train_batch_size: 1024 4 | eval_batch_size: 1024 5 | iterations_per_loop: 1000 6 | num_cores: 8 7 | skip_host_call: True 8 | augment_name: 'randaugment' 9 | randaug_num_layers: 2 10 | randaug_magnitude: 9 11 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v2-128.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 7116 2 | train_batch_size: 16384 3 | eval_batch_size: 1024 4 | iterations_per_loop: 7116 5 | skip_host_call: True 6 | num_cores: 128 7 | enable_lars: True 8 | label_smoothing: 0.1 9 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v2-256.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 3558 2 | train_batch_size: 32768 3 | eval_batch_size: 1024 4 | iterations_per_loop: 3558 5 | skip_host_call: True 6 | num_cores: 256 7 | enable_lars: True 8 | label_smoothing: 0.1 9 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v2-32.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 28464 2 | train_batch_size: 4096 3 | eval_batch_size: 1024 4 | iterations_per_loop: 28464 5 | skip_host_call: True 6 | num_cores: 32 7 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v2-512.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 3558 2 | train_batch_size: 32768 3 | eval_batch_size: 1024 4 | iterations_per_loop: 3558 5 | skip_host_call: True 6 | num_cores: 512 7 | enable_lars: True 8 | label_smoothing: 0.1 9 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v2-8.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 113854 2 | train_batch_size: 1024 3 | eval_batch_size: 1024 4 | iterations_per_loop: 113854 5 | skip_host_call: True 6 | num_cores: 8 7 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-1024.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_steps: 3558 3 | train_batch_size: 32768 4 | eval_batch_size: 1024 5 | iterations_per_loop: 3558 6 | skip_host_call: True 7 | num_cores: 1024 8 | enable_lars: True 9 | label_smoothing: 0.1 10 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-128.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_steps: 7116 3 | train_batch_size: 16384 4 | eval_batch_size: 1024 5 | iterations_per_loop: 7116 6 | skip_host_call: True 7 | num_cores: 128 8 | enable_lars: True 9 | label_smoothing: 0.1 10 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-2048.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_steps: 3558 3 | train_batch_size: 32768 4 | eval_batch_size: 1024 5 | iterations_per_loop: 3558 6 | skip_host_call: True 7 | num_cores: 2048 8 | enable_lars: True 9 | label_smoothing: 0.1 10 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-256.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_steps: 3558 3 | train_batch_size: 32768 4 | eval_batch_size: 1024 5 | iterations_per_loop: 3558 6 | skip_host_call: True 7 | num_cores: 256 8 | enable_lars: True 9 | label_smoothing: 0.1 10 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-32.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_steps: 28464 3 | train_batch_size: 4096 4 | eval_batch_size: 1024 5 | iterations_per_loop: 28464 6 | skip_host_call: True 7 | num_cores: 32 8 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-512.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_steps: 3558 3 | train_batch_size: 32768 4 | eval_batch_size: 1024 5 | iterations_per_loop: 3558 6 | skip_host_call: True 7 | num_cores: 512 8 | enable_lars: True 9 | label_smoothing: 0.1 10 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-64.yaml: -------------------------------------------------------------------------------- 1 | # DISCLAIMER: These parameters have not been optimized 2 | train_steps: 14232 3 | train_batch_size: 8192 4 | eval_batch_size: 1024 5 | iterations_per_loop: 14232 6 | skip_host_call: True 7 | num_cores: 64 8 | -------------------------------------------------------------------------------- /models/official/resnet/configs/cloud/v3-8.yaml: -------------------------------------------------------------------------------- 1 | train_steps: 113854 2 | train_batch_size: 1024 3 | eval_batch_size: 1024 4 | iterations_per_loop: 113854 5 | skip_host_call: True 6 | num_cores: 8 7 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs101_i160.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: null 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 160 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 10 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 101 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs101_i192.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: null 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 192 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 101 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs152_i192.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: null 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 192 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 152 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs152_i224.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: null 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 224 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 152 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs152_i256.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: null 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 256 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 152 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs200_i256.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: 0.1 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 256 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 200 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs270_i256.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: 0.1 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 256 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 270 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs350_i256.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: 0.1 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 256 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 350 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs350_i320.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: 0.1 5 | dropout_rate: 0.4 6 | eval_batch_size: 200 7 | image_size: 320 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 350 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs420_i320.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: 0.1 5 | dropout_rate: 0.4 6 | eval_batch_size: 200 7 | image_size: 320 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 15 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 420 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/resnet/resnet_rs/configs/resnetrs50_i160.yaml: -------------------------------------------------------------------------------- 1 | augment_name: randaugment 2 | base_learning_rate: 0.1 3 | data_format: channels_last 4 | drop_connect_rate: null 5 | dropout_rate: 0.25 6 | eval_batch_size: 200 7 | image_size: 160 8 | label_smoothing: 0.1 9 | momentum: 0.9 10 | num_eval_images: 50000 11 | num_label_classes: 1000 12 | num_train_images: 1281167 13 | precision: bfloat16 14 | randaug_magnitude: 10 15 | randaug_num_layers: 2 16 | replace_stem_max_pool: True 17 | resnet_depth: 50 18 | resnetd_shortcut: True 19 | se_ratio: 0.25 20 | train_batch_size: 4096 21 | train_steps: 109475 22 | transpose_input: True 23 | use_resnetd_stem: True 24 | use_tpu: True 25 | weight_decay: 4.0e-05 26 | moving_average_decay: 0.9999 27 | bn_momentum: 0 28 | -------------------------------------------------------------------------------- /models/official/retinanet/object_detection/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/squeezenet/configs/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /models/official/transformer/README.md: -------------------------------------------------------------------------------- 1 | Tensor2Tensor: See https://github.com/tensorflow/tensor2tensor/blob/master/docs/cloud_tpu.md 2 | 3 | BERT: See https://github.com/google-research/bert/blob/master/README.md 4 | -------------------------------------------------------------------------------- /models/official/unet3d/configs/cloud/v3-128_256x256x256_ce.yaml: -------------------------------------------------------------------------------- 1 | init_learning_rate: 0.0001 # with 0.005 the network is unstable 2 | loss: 'cross_entropy' 3 | train_batch_size: 32 4 | eval_batch_size: 1 5 | input_partition_dims: [1,16,1,1,1] 6 | use_index_label_in_train: true 7 | 8 | input_image_size: [256,256,256] 9 | label_dtype: 'float32' 10 | -------------------------------------------------------------------------------- /models/official/unet3d/configs/cloud/v3-128_256x256x256_dice.yaml: -------------------------------------------------------------------------------- 1 | init_learning_rate: 0.0001 # with 0.005 the network is unstable 2 | loss: 'adaptive_dice32' 3 | train_batch_size: 32 4 | eval_batch_size: 1 5 | input_partition_dims: [1,16,1,1,1] 6 | use_index_label_in_train: false 7 | 8 | input_image_size: [256,256,256] 9 | label_dtype: 'float32' 10 | -------------------------------------------------------------------------------- /models/official/unet3d/configs/cloud/v3-32_256x256x256_ce.yaml: -------------------------------------------------------------------------------- 1 | init_learning_rate: 0.00005 # with 0.005 the network is unstable 2 | loss: 'cross_entropy' 3 | train_batch_size: 8 4 | eval_batch_size: 1 5 | input_partition_dims: [1,16,1,1,1] 6 | use_index_label_in_train: true 7 | 8 | input_image_size: [256,256,256] 9 | label_dtype: 'float32' 10 | -------------------------------------------------------------------------------- /models/official/unet3d/configs/cloud/v3-32_256x256x256_dice.yaml: -------------------------------------------------------------------------------- 1 | init_learning_rate: 0.00005 # with 0.005 the network is unstable 2 | loss: 'adaptive_dice32' 3 | train_batch_size: 8 4 | eval_batch_size: 1 5 | input_partition_dims: [1,16,1,1,1] 6 | use_index_label_in_train: false 7 | 8 | input_image_size: [256,256,256] 9 | label_dtype: 'float32' 10 | -------------------------------------------------------------------------------- /models/official/unet3d/configs/cloud/v3-8_128x128x128_ce.yaml: -------------------------------------------------------------------------------- 1 | init_learning_rate: 0.0001 2 | loss: 'cross_entropy' 3 | train_batch_size: 16 4 | eval_batch_size: 8 5 | input_partition_dims: [1,8,1,1,1] 6 | use_index_label_in_train: true 7 | 8 | input_image_size: [128,128,128] 9 | label_dtype: 'float32' 10 | -------------------------------------------------------------------------------- /models/official/unet3d/configs/cloud/v3-8_128x128x128_dice.yaml: -------------------------------------------------------------------------------- 1 | init_learning_rate: 0.0001 # with 0.005 the network is unstable 2 | loss: 'adaptive_dice32' 3 | train_batch_size: 32 4 | eval_batch_size: 8 5 | input_partition_dims: # Yaml reads None as a string. Instead, put empty string here for NoneType. 6 | use_index_label_in_train: false 7 | 8 | input_image_size: [128,128,128] 9 | label_dtype: 'float32' 10 | -------------------------------------------------------------------------------- /models/official/unet3d/requirements.txt: -------------------------------------------------------------------------------- 1 | cloud-tpu-profiler>=1.12 2 | numpy 3 | pyyaml 4 | pytype 5 | -------------------------------------------------------------------------------- /models/setup.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Cloud TPU samples.""" 16 | 17 | from setuptools import find_packages 18 | from setuptools import setup 19 | 20 | setup( 21 | name='cloud_tpu_samples', 22 | packages=find_packages() 23 | ) 24 | -------------------------------------------------------------------------------- /tools/colab/BUILD: -------------------------------------------------------------------------------- 1 | # Colab notebooks for Cloud TPU users. 2 | 3 | package( 4 | default_applicable_licenses = ["//third_party/cloud_tpu:license"], 5 | default_visibility = ["//visibility:public"], 6 | ) 7 | 8 | licenses(["notice"]) 9 | 10 | exports_files(glob([ 11 | "*.ipynb", 12 | ])) 13 | -------------------------------------------------------------------------------- /tools/ctpu/.gitignore: -------------------------------------------------------------------------------- 1 | ctpu 2 | ctpu.exe 3 | *~ 4 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/clean/README.md: -------------------------------------------------------------------------------- 1 | # Clean # 2 | 3 | This is a complete working example, validating that everything is parsed correctly. -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/clean/active_config: -------------------------------------------------------------------------------- 1 | ctpu9 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/clean/application_default_credentials.json: -------------------------------------------------------------------------------- 1 | fake_app_default_file 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/clean/configurations/config_ctpu9: -------------------------------------------------------------------------------- 1 | [core] 2 | account = noreply@google.com 3 | unused = extraneous data 4 | project = ctpu9-test-project 5 | 6 | [compute] 7 | zone = us-central1-c 8 | region = us-central1 9 | 10 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/corrupted/README.md: -------------------------------------------------------------------------------- 1 | # Corrupted configuration environment # 2 | 3 | The gcloud active config does not exist under configurations/... 4 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/corrupted/active_config: -------------------------------------------------------------------------------- 1 | ctpu9 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/corrupted/application_default_credentials.json: -------------------------------------------------------------------------------- 1 | fake_app_default_file 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/corrupted/configurations/config_default: -------------------------------------------------------------------------------- 1 | [core] 2 | account = noreply@google.com 3 | project = ctpu9-test-project 4 | 5 | [compute] 6 | zone = us-central1-c 7 | region = us-central1 8 | 9 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/corrupted2/README.md: -------------------------------------------------------------------------------- 1 | # Corrupted 2 # 2 | 3 | There is no configurations directory. 4 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/corrupted2/active_config: -------------------------------------------------------------------------------- 1 | ctpu9 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/corrupted2/application_default_credentials.json: -------------------------------------------------------------------------------- 1 | fake_app_default_file 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/incomplete/README.md: -------------------------------------------------------------------------------- 1 | # Incomplete # 2 | 3 | This configuration simply does not have all required values set. 4 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/incomplete/active_config: -------------------------------------------------------------------------------- 1 | ctpu9 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/incomplete/application_default_credentials.json: -------------------------------------------------------------------------------- 1 | fake_app_default_file 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/incomplete/configurations/config_ctpu9: -------------------------------------------------------------------------------- 1 | [core] 2 | account = noreply@google.com 3 | 4 | [compute] 5 | zone = us-central1-c 6 | 7 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/no_app_creds/README.md: -------------------------------------------------------------------------------- 1 | # No App Creds # 2 | 3 | This is a working example, except there are no application default credentials. 4 | When running in the Cloud DevShell environment, this is a possible scenario, and 5 | should be validated to parse correctly. 6 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/no_app_creds/active_config: -------------------------------------------------------------------------------- 1 | ctpu9 2 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/no_app_creds/configurations/config_ctpu9: -------------------------------------------------------------------------------- 1 | [core] 2 | account = noreply@google.com 3 | unused = extraneous data 4 | project = ctpu9-test-project 5 | 6 | [compute] 7 | zone = us-central1-c 8 | region = us-central1 9 | 10 | -------------------------------------------------------------------------------- /tools/ctpu/config/testdata/gcloud/no_config/README.md: -------------------------------------------------------------------------------- 1 | # No configuration # 2 | 3 | This scenario is when the user has no configuration at all for gcloud. -------------------------------------------------------------------------------- /tools/data_converter/README.md: -------------------------------------------------------------------------------- 1 | # About 2 | This folder contains a suite of tools that builds upon [tensorflow/datasets](https://www.tensorflow.org/datasets) 3 | that can be used to easily convert raw data into the TFRecord format on GCS. 4 | This is helpful because data must be stored in [TFRecords](https://www.tensorflow.org/tutorials/load_data/tf_records) 5 | on [GCS](https://cloud.google.com/storage/) to run with TPU models. 6 | 7 | # High-Level Overview 8 | The folder is divided by task and each task has specific fields that are required 9 | "essential inputs" for each task. 10 | 11 | For example, image classification requires an image and a label. However, models 12 | may require more features, and this tool both facilitates the extraction of 13 | these extra features and converts the data into TFRecords. 14 | 15 | Currently supported tasks: 16 | - Image Classification 17 | 18 | # Usage 19 | To use the tool, create an implementation of one of the abstract BuilderConfigs. 20 | 21 | For example: 22 | ``` 23 | class MyBuilderConfig(ImageClassificationDataConfig): 24 | ... 25 | 26 | config = MyBuilderConfig(name="MyBuilderConfig", 27 | description="MyBuilderConfig") 28 | ds = ImageClassificationData(config) 29 | ds.download_and_prepare() 30 | 31 | ``` 32 | 33 | In each folder are also simple examples for further reference. 34 | -------------------------------------------------------------------------------- /tools/data_converter/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tools/data_converter/image_classification/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tools/data_converter/object_detection/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tools/dataset_profiler/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tools/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tools/docker/Dockerfile: -------------------------------------------------------------------------------- 1 | # Docker image for running TPU tensorflow examples. 2 | FROM ubuntu:bionic 3 | 4 | RUN apt-get update && apt-get install -y --no-install-recommends \ 5 | curl \ 6 | wget \ 7 | sudo \ 8 | gnupg \ 9 | lsb-release \ 10 | ca-certificates \ 11 | build-essential \ 12 | git \ 13 | python3 \ 14 | python3-pip \ 15 | python-setuptools && \ 16 | export CLOUD_SDK_REPO="cloud-sdk-$(lsb_release -c -s)" && \ 17 | echo "deb https://packages.cloud.google.com/apt $CLOUD_SDK_REPO main" > /etc/apt/sources.list.d/google-cloud-sdk.list && \ 18 | curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \ 19 | apt-get update && \ 20 | apt-get install -y google-cloud-sdk 21 | 22 | RUN curl https://bootstrap.pypa.io/get-pip.py -o /tmp/get-pip.py && \ 23 | python3 /tmp/get-pip.py 24 | 25 | RUN pip3 install pyyaml && \ 26 | pip3 install wheel && \ 27 | pip3 install tensorflow==2.2.0 && \ 28 | pip3 install google-cloud-storage && \ 29 | pip3 install google-api-python-client && \ 30 | pip3 install oauth2client 31 | 32 | # Checkout tensorflow/models at the appropriate branch 33 | RUN git clone -b r2.2.0 --depth 1 https://github.com/tensorflow/models.git /tensorflow_models 34 | 35 | # Checkout tensorflow/tpu at the appropriate branch 36 | RUN git clone -b r2.2.0 --depth 1 https://github.com/tensorflow/tpu.git /tensorflow_tpu_models 37 | -------------------------------------------------------------------------------- /tools/docker/Dockerfile.util: -------------------------------------------------------------------------------- 1 | # Docker image of TensorBoard and TPU Profiler. 2 | FROM ubuntu:bionic 3 | RUN apt-get update && apt-get install -y --no-install-recommends \ 4 | ca-certificates \ 5 | build-essential \ 6 | git \ 7 | python \ 8 | python-pip \ 9 | python-setuptools && \ 10 | pip install tensorflow==1.11 && \ 11 | pip install google-cloud-storage && \ 12 | pip install google-api-python-client && \ 13 | pip install oauth2client && \ 14 | pip install cloud-tpu-profiler==1.11 -------------------------------------------------------------------------------- /tools/driver/Makefile.tpu_kernel_modules: -------------------------------------------------------------------------------- 1 | KERNELDIR ?= /lib/modules/$(shell uname -r)/build 2 | PWD = $(shell pwd) 3 | EXTRA_CFLAGS = -I$(PWD)/include -I$(PWD)/include/uapi -I$(PWD) 4 | 5 | export CONFIG_HW_ACCEL = m 6 | export CONFIG_GASKET_FRAMEWORK = m 7 | export CONFIG_TPU_COMMON_COMMON = m 8 | export CONFIG_TPU_V2_CORE = m 9 | export CONFIG_TPU_V3_DRIVER = m 10 | export CONFIG_TPU_V2_DRIVER = m 11 | export CONFIG_ASIC_SW = m 12 | export CONFIG_TPU_V4_COMMON_COMMON = m 13 | export CONFIG_TPU_V4_DRIVER = m 14 | export CONFIG_TPU_V4_LITE_DRIVER = m 15 | 16 | all: 17 | $(MAKE) -C $(KERNELDIR) M=$(PWD)/drivers/accel \ 18 | EXTRA_CFLAGS="$(EXTRA_CFLAGS)" modules 19 | $(MAKE) -C $(KERNELDIR) M=$(PWD)/drivers/gasket \ 20 | EXTRA_CFLAGS="$(EXTRA_CFLAGS)" \ 21 | KBUILD_EXTRA_SYMBOLS="$(PWD)/drivers/accel/Module.symvers" \ 22 | modules 23 | $(MAKE) -C $(KERNELDIR) M=$(PWD)/drivers/asic_sw \ 24 | EXTRA_CFLAGS="$(EXTRA_CFLAGS)" \ 25 | KBUILD_EXTRA_SYMBOLS="$(PWD)/drivers/gasket/Module.symvers" \ 26 | modules 27 | $(MAKE) -C $(KERNELDIR) M=$(PWD)/drivers/char/tpu_common \ 28 | EXTRA_CFLAGS="$(EXTRA_CFLAGS)" \ 29 | KBUILD_EXTRA_SYMBOLS="$(PWD)/drivers/gasket/Module.symvers $(PWD)/drivers/asic_sw/Module.symvers" \ 30 | modules 31 | find . -name "*.ko" -exec cp -f '{}' . ';' 32 | 33 | clean: 34 | $(MAKE) -C $(KERNELDIR) M=$(PWD)/drivers/accel clean 35 | $(MAKE) -C $(KERNELDIR) M=$(PWD)/drivers/gasket clean 36 | $(MAKE) -C $(KERNELDIR) M=$(PWD)/drivers/char/tpu_common clean -------------------------------------------------------------------------------- /tools/driver/README: -------------------------------------------------------------------------------- 1 | TPU Driver. 2 | 3 | This driver as well as the TPU device are in Early Access stage [1], 4 | the feature is available to a closed group of testers. 5 | 6 | [1] https://cloud.google.com/terms/launch-stages 7 | 8 | -------------------------------------------------------------------------------- /tools/driver/drivers/accel/Kconfig: -------------------------------------------------------------------------------- 1 | # 2 | # HW Accelerator device configuration 3 | # 4 | 5 | config HW_ACCEL 6 | tristate "Google HW Accelerator Class Support" 7 | help 8 | HW Accelerator class support for Google. This option selects 9 | the framework / common code needed for the Google's hardware 10 | accelerator projects. 11 | -------------------------------------------------------------------------------- /tools/driver/drivers/accel/Makefile: -------------------------------------------------------------------------------- 1 | # 2 | # Makefile for the kernel HW accelerators device drivers 3 | # 4 | 5 | accel-class-y += accel.o 6 | accel-class-y += accel_lib.o 7 | obj-$(CONFIG_HW_ACCEL) += accel-class.o 8 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/Kconfig: -------------------------------------------------------------------------------- 1 | config ASIC_SW 2 | tristate "ASIC-SW" 3 | depends on GASKET_FRAMEWORK 4 | help 5 | A common library used by asic_sw developed kernel drivers. 6 | 7 | To compile this driver as a module, choose M here. The module 8 | will be called asic_sw. 9 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/Makefile: -------------------------------------------------------------------------------- 1 | # Makefile for source common to ASIC-SW drivers. 2 | 3 | subdir-ccflags-y := -I$(srctree) 4 | 5 | obj-$(CONFIG_ASIC_SW) += asic_sw.o 6 | 7 | asic_sw-objs := asic_sw_module.o asic_sw_firmware_indirect_registers.o \ 8 | asic_sw_clock.o 9 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_fw_device_owner_accessor.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_ASIC_SW_ASIC_FW_DEVICE_OWNER_ACCESSOR_H_ 6 | #define _DRIVERS_ASIC_SW_ASIC_FW_DEVICE_OWNER_ACCESSOR_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | static inline uint64_t asic_fw_device_owner_value(const uint64_t reg_value) 9 | { 10 | return (uint64_t)((((reg_value >> 0) & 0xffffffffffffffffULL) << 0)); 11 | } 12 | static inline int set_asic_fw_device_owner_value(uint64_t *reg_value, 13 | uint64_t value) 14 | { 15 | if ((uint64_t)value < 0x0ULL || (uint64_t)value > 0xffffffffffffffffULL) 16 | return 1; 17 | (*reg_value) = ((*reg_value) & ~((0xffffffffffffffffULL) << 0)) | 18 | (((value >> 0) & (0xffffffffffffffffULL)) << 0); 19 | return 0; 20 | } 21 | #endif 22 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_fw_device_owner_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_ASIC_SW_ASIC_FW_DEVICE_OWNER_OFFSETS_H_ 6 | #define _DRIVERS_ASIC_SW_ASIC_FW_DEVICE_OWNER_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct asic_fw_device_owner_offsets { 9 | uint64 device_owner; 10 | }; 11 | #endif 12 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_fw_host_reset_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_ASIC_SW_ASIC_FW_HOST_RESET_OFFSETS_H_ 6 | #define _DRIVERS_ASIC_SW_ASIC_FW_HOST_RESET_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct asic_fw_host_reset_offsets { 9 | uint64 chip_reset; 10 | }; 11 | #endif 12 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_fw_indirect_register_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_ASIC_SW_ASIC_FW_INDIRECT_REGISTER_OFFSETS_H_ 6 | #define _DRIVERS_ASIC_SW_ASIC_FW_INDIRECT_REGISTER_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct asic_fw_indirect_register_offsets { 9 | uint64 indirect_accessor_address; 10 | uint64 indirect_accessor_control; 11 | uint64 indirect_accessor_status; 12 | uint64 indirect_accessor_value; 13 | uint64 indirect_accessor_version; 14 | }; 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_fw_reinit_reset_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_ASIC_SW_ASIC_FW_REINIT_RESET_OFFSETS_H_ 6 | #define _DRIVERS_ASIC_SW_ASIC_FW_REINIT_RESET_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct asic_fw_reinit_reset_offsets { 9 | uint64 chip_reset; 10 | }; 11 | #endif 12 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_fw_state_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_ASIC_SW_ASIC_FW_STATE_OFFSETS_H_ 6 | #define _DRIVERS_ASIC_SW_ASIC_FW_STATE_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct asic_fw_state_offsets { 9 | uint64 firmware_status; 10 | }; 11 | #endif 12 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_fw_version_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_ASIC_SW_ASIC_FW_VERSION_OFFSETS_H_ 6 | #define _DRIVERS_ASIC_SW_ASIC_FW_VERSION_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct asic_fw_version_offsets { 9 | uint64 changelist; 10 | uint64 image_info; 11 | uint64 primary_version; 12 | uint64 secondary_version; 13 | }; 14 | #endif 15 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_sw_clock.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "drivers/asic_sw/asic_sw_clock.h" 6 | #include 7 | #include 8 | #include 9 | void asic_sw_sleep_for_msecs(uint64 duration_msecs) 10 | { 11 | set_current_state(TASK_UNINTERRUPTIBLE); 12 | schedule_timeout(msecs_to_jiffies(duration_msecs)); 13 | } 14 | EXPORT_SYMBOL(asic_sw_sleep_for_msecs); 15 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_sw_clock.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef PLATFORMS_ASIC_SW_KERNEL_COMMON_ASIC_SW_CLOCK_H_ 6 | #define PLATFORMS_ASIC_SW_KERNEL_COMMON_ASIC_SW_CLOCK_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | void asic_sw_sleep_for_msecs(uint64 duration_msecs); 9 | #endif 10 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_sw_firmware_indirect_registers.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef PLATFORMS_ASIC_SW_KERNEL_COMMON_ASIC_SW_FIRMWARE_INDIRECT_REGISTERS_H_ 6 | #define PLATFORMS_ASIC_SW_KERNEL_COMMON_ASIC_SW_FIRMWARE_INDIRECT_REGISTERS_H_ 7 | #include "drivers/asic_sw/asic_fw_indirect_register_offsets.h" 8 | #include "drivers/gasket/gasket_core.h" 9 | #include "drivers/gasket/gasket_types.h" 10 | struct firmware_indirect { 11 | const struct asic_fw_indirect_register_offsets *offsets; 12 | int bar; 13 | }; 14 | int asic_sw_firmware_indirect_read_64( 15 | struct gasket_dev *gasket_dev, 16 | const struct firmware_indirect *firmware_indirect, uint64 address, 17 | uint64 *read_value); 18 | int asic_sw_firmware_indirect_write_64( 19 | struct gasket_dev *gasket_dev, 20 | const struct firmware_indirect *firmware_indirect, uint64 address, 21 | uint64 write_value); 22 | #endif 23 | -------------------------------------------------------------------------------- /tools/driver/drivers/asic_sw/asic_sw_module.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include 6 | MODULE_DESCRIPTION("ASIC-SW Driver Library"); 7 | MODULE_VERSION("1.0.0.0"); 8 | MODULE_LICENSE("GPL"); 9 | MODULE_AUTHOR("ASIC-SW "); 10 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_H_ 7 | #define TPUV4_LBUS_BAR_INDEX 0x0 8 | #define TPUV4_LBUS_BAR_OFFSET 0x0 9 | #define TPUV4_LBUS_BAR_BYTES 0x4000000 10 | #define TPUV4_CSR_BAR_INDEX 0x2 11 | #define TPUV4_CSR_BAR_OFFSET 0x10000000 12 | #define TPUV4_CSR_BAR_BYTES 0x80000000 13 | #define TPUV4_RESET_RETRY 1200 14 | #define TPUV4_RESET_DELAY 10 15 | #define TPUV4_CHIP_INIT_DONE 1 16 | #define TPUV4_RESET_ACCEPTED 0 17 | enum tpuv4_reset_types { TPUV4_CHIP_REINIT_RESET = 1 }; 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_bar0_ranges.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_bar0_ranges.h" 6 | static const struct gasket_mappable_region 7 | tpu_common_security_level_root_ranges[] = { 8 | { 0x0, 0x250000, VM_READ | VM_WRITE }, 9 | { 0x260000, 0x30000, VM_READ | VM_WRITE }, 10 | { 0x2a0000, 0x50000, VM_READ | VM_WRITE }, 11 | { 0x400000, 0x40000, VM_READ | VM_WRITE }, 12 | }; 13 | static const struct gasket_mappable_region 14 | tpu_common_security_level_user_ranges[] = { 15 | { 0x0, 0x250000, VM_READ }, 16 | { 0x260000, 0x30000, VM_READ }, 17 | { 0x2a0000, 0x50000, VM_READ }, 18 | { 0x400000, 0x40000, VM_READ }, 19 | }; 20 | int tpuv4_bar0_get_region_count(enum tpu_common_security_level group) 21 | { 22 | if (group == TPU_COMMON_SECURITY_LEVEL_ROOT) 23 | return ARRAY_SIZE(tpu_common_security_level_root_ranges); 24 | else if (group == TPU_COMMON_SECURITY_LEVEL_USER) 25 | return ARRAY_SIZE(tpu_common_security_level_user_ranges); 26 | else 27 | return 0; 28 | } 29 | const struct gasket_mappable_region * 30 | tpuv4_bar0_get_regions(enum tpu_common_security_level group) 31 | { 32 | if (group == TPU_COMMON_SECURITY_LEVEL_ROOT) 33 | return tpu_common_security_level_root_ranges; 34 | else if (group == TPU_COMMON_SECURITY_LEVEL_USER) 35 | return tpu_common_security_level_user_ranges; 36 | else 37 | return NULL; 38 | } 39 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_bar0_ranges.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_BAR0_RANGES_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_BAR0_RANGES_H_ 7 | #include "drivers/char/tpu_common/tpu_common.h" 8 | #include "drivers/gasket/gasket_core.h" 9 | #include 10 | #include 11 | #include 12 | #if LINUX_VERSION_CODE < KERNEL_VERSION(5, 15, 0) 13 | #include 14 | #endif 15 | int tpuv4_bar0_get_region_count(enum tpu_common_security_level group); 16 | const struct gasket_mappable_region * 17 | tpuv4_bar0_get_regions(enum tpu_common_security_level group); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_bar2_ranges.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_BAR2_RANGES_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_BAR2_RANGES_H_ 7 | #include "drivers/char/tpu_common/tpu_common.h" 8 | #include "drivers/gasket/gasket_core.h" 9 | #include 10 | #include 11 | #include 12 | #if LINUX_VERSION_CODE < KERNEL_VERSION(5, 15, 0) 13 | #include 14 | #endif 15 | int tpuv4_bar2_get_region_count(enum tpu_common_security_level group); 16 | const struct gasket_mappable_region * 17 | tpuv4_bar2_get_regions(enum tpu_common_security_level group); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_device_owner_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_device_owner_offsets.h" 6 | static const struct asic_fw_device_owner_offsets 7 | tpuv4_device_owner_offsets_all_offsets[1] = { 8 | { 0x4720010 } 9 | }; 10 | int tpuv4_device_owner_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct asic_fw_device_owner_offsets * 15 | tpuv4_device_owner_offsets_get(int index) 16 | { 17 | if (index < 0 || index >= 1) 18 | return NULL; 19 | return &tpuv4_device_owner_offsets_all_offsets[index]; 20 | } 21 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_device_owner_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_DEVICE_OWNER_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_DEVICE_OWNER_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_device_owner_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_device_owner_offsets { 10 | uint64 device_owner; 11 | }; 12 | int tpuv4_device_owner_offsets_count(void); 13 | const struct asic_fw_device_owner_offsets * 14 | tpuv4_device_owner_offsets_get(int index); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_firmware_state_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_firmware_state_offsets.h" 6 | static const struct asic_fw_state_offsets 7 | tpuv4_firmware_state_offsets_all_offsets[1] = { 8 | { 0x46d0000 } 9 | }; 10 | int tpuv4_firmware_state_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct asic_fw_state_offsets *tpuv4_firmware_state_offsets_get(int index) 15 | { 16 | if (index < 0 || index >= 1) 17 | return NULL; 18 | return &tpuv4_firmware_state_offsets_all_offsets[index]; 19 | } 20 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_firmware_state_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_STATE_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_STATE_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_state_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_firmware_state_offsets { 10 | uint64 firmware_status; 11 | }; 12 | int tpuv4_firmware_state_offsets_count(void); 13 | const struct asic_fw_state_offsets * 14 | tpuv4_firmware_state_offsets_get(int index); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_firmware_version_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_firmware_version_offsets.h" 6 | static const struct asic_fw_version_offsets 7 | tpuv4_firmware_version_offsets_all_offsets[1] = { 8 | { 0x46d0050, 0x46d0048, 9 | 0x46d0038, 10 | 0x46d0040 } 11 | }; 12 | int tpuv4_firmware_version_offsets_count(void) 13 | { 14 | return 1; 15 | } 16 | const struct asic_fw_version_offsets * 17 | tpuv4_firmware_version_offsets_get(int index) 18 | { 19 | if (index < 0 || index >= 1) 20 | return NULL; 21 | return &tpuv4_firmware_version_offsets_all_offsets[index]; 22 | } 23 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_firmware_version_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_VERSION_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_VERSION_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_version_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_firmware_version_offsets { 10 | uint64 changelist; 11 | uint64 image_info; 12 | uint64 primary_version; 13 | uint64 secondary_version; 14 | }; 15 | int tpuv4_firmware_version_offsets_count(void); 16 | const struct asic_fw_version_offsets * 17 | tpuv4_firmware_version_offsets_get(int index); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_mgr_error_loperf_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_mgr_error_loperf_offsets.h" 6 | static const struct tpuv4common_mgr_error_loperf_offsets 7 | tpuv4_mgr_error_loperf_offsets_all_offsets[1] = { 8 | { 0x46e0018, 9 | 0x46e0000, 10 | 0x46e0008, 11 | 0x46e0010 } 12 | }; 13 | int tpuv4_mgr_error_loperf_offsets_count(void) 14 | { 15 | return 1; 16 | } 17 | const struct tpuv4common_mgr_error_loperf_offsets * 18 | tpuv4_mgr_error_loperf_offsets_get(int index) 19 | { 20 | if (index < 0 || index >= 1) 21 | return NULL; 22 | return &tpuv4_mgr_error_loperf_offsets_all_offsets[index]; 23 | } 24 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_mgr_error_loperf_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPU_V4_COMMON_MGR_ERROR_LOPERF_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPU_V4_COMMON_MGR_ERROR_LOPERF_OFFSETS_H_ 7 | #include "drivers/char/tpu_common/tpuv4common/tpuv4common_mgr_error_loperf_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_mgr_error_loperf_offsets { 10 | uint64 bootloader_running_timer; 11 | uint64 error_interrupt_control; 12 | uint64 error_interrupt_status; 13 | uint64 global_fatal_error_status; 14 | }; 15 | int tpuv4_mgr_error_loperf_offsets_count(void); 16 | const struct tpuv4common_mgr_error_loperf_offsets * 17 | tpuv4_mgr_error_loperf_offsets_get(int index); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_page_table_config.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_PAGE_TABLE_CONFIG_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPUV4_PAGE_TABLE_CONFIG_H_ 7 | #include "drivers/gasket/gasket_core.h" 8 | #include 9 | #include 10 | const struct gasket_page_table_config tpuv4_page_table_configs[] = { 11 | { 0, GASKET_PAGE_TABLE_MODE_EXTENDED, 0x20000, 2, 0x13b0000, 0, 0 }, 12 | }; 13 | #endif 14 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_pcie_flr_status_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_pcie_flr_status_offsets.h" 6 | static const struct tpuv4common_pcie_flr_status_offsets 7 | tpuv4_pcie_flr_status_offsets_all_offsets[1] = { 8 | { 0x4142000 } 9 | }; 10 | int tpuv4_pcie_flr_status_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct tpuv4common_pcie_flr_status_offsets * 15 | tpuv4_pcie_flr_status_offsets_get(int index) 16 | { 17 | if (index < 0 || index >= 1) 18 | return NULL; 19 | return &tpuv4_pcie_flr_status_offsets_all_offsets[index]; 20 | } 21 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_pcie_flr_status_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPU_V4_COMMON_PCIE_FLR_STATUS_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPU_V4_COMMON_PCIE_FLR_STATUS_OFFSETS_H_ 7 | #include "drivers/char/tpu_common/tpuv4common/tpuv4common_pcie_flr_status_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_pcie_flr_status_offsets { 10 | uint64 pcie_flr_status; 11 | }; 12 | int tpuv4_pcie_flr_status_offsets_count(void); 13 | const struct tpuv4common_pcie_flr_status_offsets * 14 | tpuv4_pcie_flr_status_offsets_get(int index); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_reinit_reset_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_reinit_reset_offsets.h" 6 | static const struct asic_fw_reinit_reset_offsets 7 | tpuv4_reinit_reset_offsets_all_offsets[1] = { 8 | { 0x4720008 } 9 | }; 10 | int tpuv4_reinit_reset_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct asic_fw_reinit_reset_offsets * 15 | tpuv4_reinit_reset_offsets_get(int index) 16 | { 17 | if (index < 0 || index >= 1) 18 | return NULL; 19 | return &tpuv4_reinit_reset_offsets_all_offsets[index]; 20 | } 21 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_reinit_reset_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_REINIT_RESET_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_REINIT_RESET_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_reinit_reset_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_reinit_reset_offsets { 10 | uint64 chip_reset; 11 | }; 12 | int tpuv4_reinit_reset_offsets_count(void); 13 | const struct asic_fw_reinit_reset_offsets * 14 | tpuv4_reinit_reset_offsets_get(int index); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_userspace_firmware_indirect_accessor_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4_userspace_firmware_indirect_accessor_offsets.h" 6 | static const struct asic_fw_indirect_register_offsets 7 | tpuv4_userspace_firmware_indirect_accessor_offsets_all_offsets[1] = { 8 | { 0x46d0018, 9 | 0x46d0020, 10 | 0x46d0028, 11 | 0x46d0030, 12 | 0x46d0010 } 13 | }; 14 | int tpuv4_userspace_firmware_indirect_accessor_offsets_count(void) 15 | { 16 | return 1; 17 | } 18 | const struct asic_fw_indirect_register_offsets * 19 | tpuv4_userspace_firmware_indirect_accessor_offsets_get(int index) 20 | { 21 | if (index < 0 || index >= 1) 22 | return NULL; 23 | return &tpuv4_userspace_firmware_indirect_accessor_offsets_all_offsets 24 | [index]; 25 | } 26 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_userspace_firmware_indirect_accessor_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_INDIRECT_REGISTER_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_ASIC_FW_INDIRECT_REGISTER_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_indirect_register_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_userspace_firmware_indirect_accessor_offsets { 10 | uint64 indirect_accessor_address; 11 | uint64 indirect_accessor_control; 12 | uint64 indirect_accessor_status; 13 | uint64 indirect_accessor_value; 14 | uint64 indirect_accessor_version; 15 | }; 16 | int tpuv4_userspace_firmware_indirect_accessor_offsets_count(void); 17 | const struct asic_fw_indirect_register_offsets * 18 | tpuv4_userspace_firmware_indirect_accessor_offsets_get(int index); 19 | #endif 20 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4/tpuv4_userspace_lst_port_indirect_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPU_V4_COMMON_USERSPACE_LST_PORT_INDIRECT_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4_TPU_V4_COMMON_USERSPACE_LST_PORT_INDIRECT_OFFSETS_H_ 7 | #include "drivers/char/tpu_common/tpuv4common/tpuv4common_userspace_lst_port_indirect_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4_userspace_lst_port_indirect_offsets { 10 | uint64 data_link_layer_request; 11 | uint64 data_link_layer_status; 12 | uint64 unused_register_one; 13 | uint64 unused_register_two; 14 | uint64 unused_register_three; 15 | uint64 unused_register_four[2]; 16 | uint64 unused_register_five[2]; 17 | uint64 unused_register_six[2]; 18 | uint64 lock; 19 | uint64 physical_layer_state; 20 | uint64 rates; 21 | uint64 to_mirror[12]; 22 | }; 23 | int tpuv4_userspace_lst_port_indirect_offsets_count(void); 24 | const struct tpuv4common_userspace_lst_port_indirect_offsets * 25 | tpuv4_userspace_lst_port_indirect_offsets_get(int index); 26 | #endif 27 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4common/tpuv4common_lst.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef PLATFORMS_ASIC_SW_KERNEL_TPU_COMMON_TPU_V4_COMMON_COMMON_TPU_V4_COMMON_LST_H_ 6 | #define PLATFORMS_ASIC_SW_KERNEL_TPU_COMMON_TPU_V4_COMMON_COMMON_TPU_V4_COMMON_LST_H_ 7 | #include "drivers/asic_sw/asic_fw_indirect_register_offsets.h" 8 | #include "drivers/asic_sw/asic_sw_firmware_indirect_registers.h" 9 | #include "drivers/char/tpu_common/tpuv4common/tpuv4common_userspace_lst_port_indirect_offsets.h" 10 | #include "drivers/gasket/gasket_core.h" 11 | #include "drivers/gasket/gasket_types.h" 12 | int tpuv4common_lst_disable_lst_data_links( 13 | struct gasket_dev *gasket_dev, 14 | const struct firmware_indirect *firmware_indirect, 15 | const struct tpuv4common_userspace_lst_port_indirect_offsets 16 | *lst_userspace_port_indirect_offsets, 17 | int disable_poll_count, uint64 disable_poll_interval_msecs); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4common/tpuv4common_mgr_error_loperf_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPU_V4_COMMON_TPU_V4_COMMON_MGR_ERROR_LOPERF_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPU_V4_COMMON_TPU_V4_COMMON_MGR_ERROR_LOPERF_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct tpuv4common_mgr_error_loperf_offsets { 9 | uint64 bootloader_running_timer; 10 | uint64 error_interrupt_control; 11 | uint64 error_interrupt_status; 12 | uint64 global_fatal_error_status; 13 | }; 14 | #endif 15 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4common/tpuv4common_pcie_flr_status_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPU_V4_COMMON_TPU_V4_COMMON_PCIE_FLR_STATUS_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPU_V4_COMMON_TPU_V4_COMMON_PCIE_FLR_STATUS_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct tpuv4common_pcie_flr_status_offsets { 9 | uint64 pcie_flr_status; 10 | }; 11 | #endif 12 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4common/tpuv4common_userspace_lst_port_indirect_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPU_V4_COMMON_TPU_V4_COMMON_USERSPACE_LST_PORT_INDIRECT_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPU_V4_COMMON_TPU_V4_COMMON_USERSPACE_LST_PORT_INDIRECT_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct tpuv4common_userspace_lst_port_indirect_offsets { 9 | uint64 data_link_layer_request; 10 | uint64 data_link_layer_status; 11 | uint64 unused_register_one; 12 | uint64 unused_register_two; 13 | uint64 unused_register_three; 14 | uint64 unused_register_four[2]; 15 | uint64 unused_register_five[2]; 16 | uint64 unused_register_six[2]; 17 | uint64 lock; 18 | uint64 physical_layer_state; 19 | uint64 rates; 20 | uint64 to_mirror[12]; 21 | }; 22 | #endif 23 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_H_ 7 | #define TPUV4LITE_LBUS_BAR_INDEX 0x0 8 | #define TPUV4LITE_LBUS_BAR_OFFSET 0x0 9 | #define TPUV4LITE_LBUS_BAR_BYTES 0x4000000 10 | #define TPUV4LITE_CSR_BAR_INDEX 0x2 11 | #define TPUV4LITE_CSR_BAR_OFFSET 0x10000000 12 | #define TPUV4LITE_CSR_BAR_BYTES 0x80000000 13 | #define TPUV4LITE_RESET_RETRY 1200 14 | #define TPUV4LITE_RESET_DELAY 10 15 | #define TPUV4LITE_CHIP_INIT_DONE 1 16 | #define TPUV4LITE_RESET_ACCEPTED 0 17 | enum tpuv4lite_reset_types { TPUV4LITE_CHIP_REINIT_RESET = 1 }; 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_bar0_ranges.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_bar0_ranges.h" 6 | static const struct gasket_mappable_region 7 | tpu_common_security_level_root_ranges[] = { 8 | { 0x0, 0x250000, VM_READ | VM_WRITE }, 9 | { 0x260000, 0x30000, VM_READ | VM_WRITE }, 10 | { 0x2a0000, 0x50000, VM_READ | VM_WRITE }, 11 | { 0x400000, 0x40000, VM_READ | VM_WRITE }, 12 | }; 13 | static const struct gasket_mappable_region 14 | tpu_common_security_level_user_ranges[] = { 15 | { 0x0, 0x250000, VM_READ }, 16 | { 0x260000, 0x30000, VM_READ }, 17 | { 0x2a0000, 0x50000, VM_READ }, 18 | { 0x400000, 0x40000, VM_READ }, 19 | }; 20 | int tpuv4lite_bar0_get_region_count(enum tpu_common_security_level group) 21 | { 22 | if (group == TPU_COMMON_SECURITY_LEVEL_ROOT) 23 | return ARRAY_SIZE(tpu_common_security_level_root_ranges); 24 | else if (group == TPU_COMMON_SECURITY_LEVEL_USER) 25 | return ARRAY_SIZE(tpu_common_security_level_user_ranges); 26 | else 27 | return 0; 28 | } 29 | const struct gasket_mappable_region * 30 | tpuv4lite_bar0_get_regions(enum tpu_common_security_level group) 31 | { 32 | if (group == TPU_COMMON_SECURITY_LEVEL_ROOT) 33 | return tpu_common_security_level_root_ranges; 34 | else if (group == TPU_COMMON_SECURITY_LEVEL_USER) 35 | return tpu_common_security_level_user_ranges; 36 | else 37 | return NULL; 38 | } 39 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_bar0_ranges.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_BAR0_RANGES_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_BAR0_RANGES_H_ 7 | #include "drivers/char/tpu_common/tpu_common.h" 8 | #include "drivers/gasket/gasket_core.h" 9 | #include 10 | #include 11 | #include 12 | #if LINUX_VERSION_CODE < KERNEL_VERSION(5, 15, 0) 13 | #include 14 | #endif 15 | int tpuv4lite_bar0_get_region_count(enum tpu_common_security_level group); 16 | const struct gasket_mappable_region * 17 | tpuv4lite_bar0_get_regions(enum tpu_common_security_level group); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_bar2_ranges.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_BAR2_RANGES_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_BAR2_RANGES_H_ 7 | #include "drivers/char/tpu_common/tpu_common.h" 8 | #include "drivers/gasket/gasket_core.h" 9 | #include 10 | #include 11 | #include 12 | #if LINUX_VERSION_CODE < KERNEL_VERSION(5, 15, 0) 13 | #include 14 | #endif 15 | int tpuv4lite_bar2_get_region_count(enum tpu_common_security_level group); 16 | const struct gasket_mappable_region * 17 | tpuv4lite_bar2_get_regions(enum tpu_common_security_level group); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_device_owner_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_device_owner_offsets.h" 6 | static const struct asic_fw_device_owner_offsets 7 | tpuv4lite_device_owner_offsets_all_offsets[1] = { 8 | { 0x4720010 } 9 | }; 10 | int tpuv4lite_device_owner_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct asic_fw_device_owner_offsets * 15 | tpuv4lite_device_owner_offsets_get(int index) 16 | { 17 | if (index < 0 || index >= 1) 18 | return NULL; 19 | return &tpuv4lite_device_owner_offsets_all_offsets[index]; 20 | } 21 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_device_owner_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_DEVICE_OWNER_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_DEVICE_OWNER_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_device_owner_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4lite_device_owner_offsets { 10 | uint64 device_owner; 11 | }; 12 | int tpuv4lite_device_owner_offsets_count(void); 13 | const struct asic_fw_device_owner_offsets * 14 | tpuv4lite_device_owner_offsets_get(int index); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_firmware_state_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_firmware_state_offsets.h" 6 | static const struct asic_fw_state_offsets 7 | tpuv4lite_firmware_state_offsets_all_offsets[1] = { 8 | { 0x46d0000 } 9 | }; 10 | int tpuv4lite_firmware_state_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct asic_fw_state_offsets *tpuv4lite_firmware_state_offsets_get(int index) 15 | { 16 | if (index < 0 || index >= 1) 17 | return NULL; 18 | return &tpuv4lite_firmware_state_offsets_all_offsets[index]; 19 | } 20 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_firmware_state_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_STATE_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_STATE_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_state_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4lite_firmware_state_offsets { 10 | uint64 firmware_status; 11 | }; 12 | int tpuv4lite_firmware_state_offsets_count(void); 13 | const struct asic_fw_state_offsets *tpuv4lite_firmware_state_offsets_get(int index); 14 | #endif 15 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_firmware_version_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_firmware_version_offsets.h" 6 | static const struct tpuv4lite_firmware_version_offsets 7 | tpuv4lite_firmware_version_offsets_all_offsets[1] = { 8 | { 0x46d0050, 0x46d0048, 9 | 0x46d0038, 10 | 0x46d0040 } 11 | }; 12 | int tpuv4lite_firmware_version_offsets_count(void) 13 | { 14 | return 1; 15 | } 16 | const struct tpuv4lite_firmware_version_offsets * 17 | tpuv4lite_firmware_version_offsets_get(int index) 18 | { 19 | if (index < 0 || index >= 1) 20 | return NULL; 21 | return &tpuv4lite_firmware_version_offsets_all_offsets[index]; 22 | } 23 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_firmware_version_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_FIRMWARE_VERSION_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_FIRMWARE_VERSION_OFFSETS_H_ 7 | #include "drivers/gasket/gasket_types.h" 8 | struct tpuv4lite_firmware_version_offsets { 9 | uint64 changelist; 10 | uint64 image_info; 11 | uint64 primary_version; 12 | uint64 secondary_version; 13 | }; 14 | int tpuv4lite_firmware_version_offsets_count(void); 15 | const struct tpuv4lite_firmware_version_offsets * 16 | tpuv4lite_firmware_version_offsets_get(int index); 17 | #endif 18 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_interrupt_desc.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_INTERRUPT_DESC_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_INTERRUPT_DESC_H_ 7 | #include "drivers/gasket/gasket_core.h" 8 | #include 9 | #include 10 | const struct gasket_interrupt_desc tpuv4lite_interrupts[] = { 11 | { 2, 0x19b0020 }, 12 | { 2, 0x19b0000 }, 13 | { 2, 0x19b0008 }, 14 | { 2, 0x19b0010 }, 15 | { 2, 0x19b0018 }, 16 | { 2, 0x1ab0020 }, 17 | { 2, 0x1ab0000 }, 18 | { 2, 0x1ab0008 }, 19 | { 2, 0x1ab0010 }, 20 | { 2, 0x1ab0018 }, 21 | { 2, 0x4720000 }, 22 | { 2, 0x1bb0000 }, 23 | { 2, 0x1bb0008 }, 24 | { 2, 0x1bb0010 }, 25 | { 2, 0x1bb0018 }, 26 | { 2, 0x90000 }, 27 | { 2, 0xb0000 }, 28 | { 2, 0xd0000 }, 29 | { 2, 0xf0000 }, 30 | { 2, 0x110000 }, 31 | { 2, 0x130000 }, 32 | { 2, 0x150000 }, 33 | { 2, 0x170000 }, 34 | { 2, 0x190000 }, 35 | { 2, 0x1b0000 }, 36 | { 2, 0x1d0000 }, 37 | { 2, 0x1f0000 }, 38 | { 2, 0x210000 }, 39 | { 2, 0x230000 }, 40 | { 2, 0x250000 }, 41 | { 2, 0x270000 }, 42 | { 2, 0x290000 }, 43 | { 2, 0x2b0000 }, 44 | { 2, 0x2d0000 }, 45 | { 2, 0x2f0000 }, 46 | { 2, 0x310000 }, 47 | { 2, 0x4720018 }, 48 | }; 49 | #endif 50 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_mgr_error_loperf_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_mgr_error_loperf_offsets.h" 6 | static const struct tpuv4common_mgr_error_loperf_offsets 7 | tpuv4lite_mgr_error_loperf_offsets_all_offsets[1] = { 8 | { 0x46e0018, 9 | 0x46e0000, 10 | 0x46e0008, 11 | 0x46e0010 } 12 | }; 13 | int tpuv4lite_mgr_error_loperf_offsets_count(void) 14 | { 15 | return 1; 16 | } 17 | const struct tpuv4common_mgr_error_loperf_offsets * 18 | tpuv4lite_mgr_error_loperf_offsets_get(int index) 19 | { 20 | if (index < 0 || index >= 1) 21 | return NULL; 22 | return &tpuv4lite_mgr_error_loperf_offsets_all_offsets[index]; 23 | } 24 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_mgr_error_loperf_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPU_V4_COMMON_MGR_ERROR_LOPERF_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPU_V4_COMMON_MGR_ERROR_LOPERF_OFFSETS_H_ 7 | #include "drivers/char/tpu_common/tpuv4common/tpuv4common_mgr_error_loperf_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4lite_mgr_error_loperf_offsets { 10 | uint64 bootloader_running_timer; 11 | uint64 error_interrupt_control; 12 | uint64 error_interrupt_status; 13 | uint64 global_fatal_error_status; 14 | }; 15 | int tpuv4lite_mgr_error_loperf_offsets_count(void); 16 | const struct tpuv4common_mgr_error_loperf_offsets * 17 | tpuv4lite_mgr_error_loperf_offsets_get(int index); 18 | #endif 19 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_page_table_config.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_PAGE_TABLE_CONFIG_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPUV4LITE_PAGE_TABLE_CONFIG_H_ 7 | #include "drivers/gasket/gasket_core.h" 8 | #include 9 | #include 10 | const struct gasket_page_table_config tpuv4lite_page_table_configs[] = { 11 | { 0, GASKET_PAGE_TABLE_MODE_EXTENDED, 0x20000, 2, 0x13b0000, 0, 0 }, 12 | }; 13 | #endif 14 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_pcie_flr_status_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_pcie_flr_status_offsets.h" 6 | static const struct tpuv4common_pcie_flr_status_offsets 7 | tpuv4lite_pcie_flr_status_offsets_all_offsets[1] = { 8 | { 0x4142000 } 9 | }; 10 | int tpuv4lite_pcie_flr_status_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct tpuv4common_pcie_flr_status_offsets * 15 | tpuv4lite_pcie_flr_status_offsets_get(int index) 16 | { 17 | if (index < 0 || index >= 1) 18 | return NULL; 19 | return &tpuv4lite_pcie_flr_status_offsets_all_offsets[index]; 20 | } 21 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_pcie_flr_status_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPU_V4_COMMON_PCIE_FLR_STATUS_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPU_V4_COMMON_PCIE_FLR_STATUS_OFFSETS_H_ 7 | #include "drivers/char/tpu_common/tpuv4common/tpuv4common_pcie_flr_status_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4lite_pcie_flr_status_offsets { 10 | uint64 pcie_flr_status; 11 | }; 12 | int tpuv4lite_pcie_flr_status_offsets_count(void); 13 | const struct tpuv4common_pcie_flr_status_offsets * 14 | tpuv4lite_pcie_flr_status_offsets_get(int index); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_reinit_reset_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_reinit_reset_offsets.h" 6 | static const struct asic_fw_reinit_reset_offsets 7 | tpuv4lite_reinit_reset_offsets_all_offsets[1] = { 8 | { 0x4720008 } 9 | }; 10 | int tpuv4lite_reinit_reset_offsets_count(void) 11 | { 12 | return 1; 13 | } 14 | const struct asic_fw_reinit_reset_offsets * 15 | tpuv4lite_reinit_reset_offsets_get(int index) 16 | { 17 | if (index < 0 || index >= 1) 18 | return NULL; 19 | return &tpuv4lite_reinit_reset_offsets_all_offsets[index]; 20 | } 21 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_reinit_reset_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_REINIT_RESET_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_REINIT_RESET_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_reinit_reset_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4lite_reinit_reset_offsets { 10 | uint64 chip_reset; 11 | }; 12 | int tpuv4lite_reinit_reset_offsets_count(void); 13 | const struct asic_fw_reinit_reset_offsets * 14 | tpuv4lite_reinit_reset_offsets_get(int index); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_userspace_firmware_indirect_accessor_offsets.c: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #include "tpuv4lite_userspace_firmware_indirect_accessor_offsets.h" 6 | static const struct asic_fw_indirect_register_offsets 7 | tpuv4lite_userspace_firmware_indirect_accessor_offsets_all_offsets[1] = { 8 | { 0x46d0018, 9 | 0x46d0020, 10 | 0x46d0028, 11 | 0x46d0030, 12 | 0x46d0010 } 13 | }; 14 | int tpuv4lite_userspace_firmware_indirect_accessor_offsets_count(void) 15 | { 16 | return 1; 17 | } 18 | const struct asic_fw_indirect_register_offsets * 19 | tpuv4lite_userspace_firmware_indirect_accessor_offsets_get(int index) 20 | { 21 | if (index < 0 || index >= 1) 22 | return NULL; 23 | return &tpuv4lite_userspace_firmware_indirect_accessor_offsets_all_offsets 24 | [index]; 25 | } 26 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_userspace_firmware_indirect_accessor_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_INDIRECT_REGISTER_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_ASIC_FW_INDIRECT_REGISTER_OFFSETS_H_ 7 | #include "drivers/asic_sw/asic_fw_indirect_register_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4lite_userspace_firmware_indirect_accessor_offsets { 10 | uint64 indirect_accessor_address; 11 | uint64 indirect_accessor_control; 12 | uint64 indirect_accessor_status; 13 | uint64 indirect_accessor_value; 14 | uint64 indirect_accessor_version; 15 | }; 16 | int tpuv4lite_userspace_firmware_indirect_accessor_offsets_count(void); 17 | const struct asic_fw_indirect_register_offsets * 18 | tpuv4lite_userspace_firmware_indirect_accessor_offsets_get(int index); 19 | #endif 20 | -------------------------------------------------------------------------------- /tools/driver/drivers/char/tpu_common/tpuv4lite/tpuv4lite_userspace_lst_port_indirect_offsets.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPU_V4_COMMON_USERSPACE_LST_PORT_INDIRECT_OFFSETS_H_ 6 | #define _DRIVERS_CHAR_TPU_COMMON_TPUV4LITE_TPU_V4_COMMON_USERSPACE_LST_PORT_INDIRECT_OFFSETS_H_ 7 | #include "drivers/char/tpu_common/tpuv4common/tpuv4common_userspace_lst_port_indirect_offsets.h" 8 | #include "drivers/gasket/gasket_types.h" 9 | struct tpuv4lite_userspace_lst_port_indirect_offsets { 10 | uint64 data_link_layer_request; 11 | uint64 data_link_layer_status; 12 | uint64 unused_register_one; 13 | uint64 unused_register_two; 14 | uint64 unused_register_three; 15 | uint64 unused_register_four[2]; 16 | uint64 unused_register_five[2]; 17 | uint64 unused_register_six[2]; 18 | uint64 lock; 19 | uint64 physical_layer_state; 20 | uint64 rates; 21 | uint64 to_mirror[12]; 22 | }; 23 | int tpuv4lite_userspace_lst_port_indirect_offsets_count(void); 24 | const struct tpuv4common_userspace_lst_port_indirect_offsets * 25 | tpuv4lite_userspace_lst_port_indirect_offsets_get(int index); 26 | #endif 27 | -------------------------------------------------------------------------------- /tools/driver/drivers/gasket/Kconfig: -------------------------------------------------------------------------------- 1 | config GASKET_FRAMEWORK 2 | select DMA_SHARED_BUFFER 3 | tristate "Gasket framework" 4 | depends on PCI 5 | depends on HW_ACCEL 6 | help 7 | This is a framework for lightweight kernel drivers, providing simplified 8 | support for MSI-x interrupts, two-level page tables, sysfs nodes, on 9 | character devices over PCI. 10 | 11 | To compile this driver as a module, choose M here. The module 12 | will be called gasket. 13 | -------------------------------------------------------------------------------- /tools/driver/drivers/gasket/Makefile: -------------------------------------------------------------------------------- 1 | # 2 | # Makefile for Gasket framework. 3 | # 4 | 5 | obj-$(CONFIG_GASKET_FRAMEWORK) += gasket.o 6 | 7 | gasket-objs := gasket_core.o gasket_dmabuf.o gasket_interrupt.o gasket_ioctl.o gasket_page_table.o gasket_sysfs.o 8 | -------------------------------------------------------------------------------- /tools/driver/drivers/gasket/README: -------------------------------------------------------------------------------- 1 | This directory contains the Gasket thin driver framework. Drivers using it may 2 | be found in other subdirectories. 3 | -------------------------------------------------------------------------------- /tools/driver/drivers/gasket/gasket_constants.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef __GASKET_CONSTANTS_H__ 6 | #define __GASKET_CONSTANTS_H__ 7 | #define GASKET_FRAMEWORK_VERSION "1.1.1" 8 | #define GASKET_FRAMEWORK_DESC_MAX 4 9 | #define GASKET_DEV_MAX 256 10 | #define GASKET_NUM_BARS 6 11 | #define GASKET_MAX_NUM_PAGE_TABLES 128 12 | #define GASKET_NAME_MAX 32 13 | enum gasket_status { 14 | GASKET_STATUS_DEAD = 0, 15 | GASKET_STATUS_LAMED, 16 | GASKET_STATUS_ALIVE, 17 | GASKET_STATUS_DRIVER_EXIT, 18 | }; 19 | #endif 20 | -------------------------------------------------------------------------------- /tools/driver/drivers/gasket/gasket_dmabuf.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef __GASKET_DMABUF_H__ 6 | #define __GASKET_DMABUF_H__ 7 | #include "gasket_types.h" 8 | struct gasket_dma_buf_device_data; 9 | struct gasket_dma_buf_device_data { 10 | void (*release_cb)(struct gasket_dma_buf_device_data *device_data); 11 | }; 12 | struct dma_buf *gasket_create_mmap_dma_buf(struct gasket_dev *gasket_dev, 13 | u64 mmap_offset, size_t buf_size, int flags, 14 | struct gasket_dma_buf_device_data *device_data); 15 | #endif 16 | -------------------------------------------------------------------------------- /tools/driver/drivers/gasket/gasket_ioctl.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef __GASKET_IOCTL_H__ 6 | #define __GASKET_IOCTL_H__ 7 | #include "gasket_core.h" 8 | long gasket_handle_ioctl(struct file *filp, uint cmd, ulong arg); 9 | long gasket_is_supported_ioctl(uint cmd); 10 | #endif 11 | -------------------------------------------------------------------------------- /tools/driver/drivers/gasket/gasket_types.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef PLATFORMS_ASIC_SW_KERNEL_COMMON_GASKET_GASKET_TYPES_H_ 6 | #define PLATFORMS_ASIC_SW_KERNEL_COMMON_GASKET_GASKET_TYPES_H_ 7 | #include 8 | typedef uint8_t uint8; 9 | typedef uint16_t uint16; 10 | typedef uint32_t uint32; 11 | typedef unsigned long long uint64; 12 | typedef int8_t int8; 13 | typedef int16_t int16; 14 | typedef int32_t int32; 15 | typedef long long int64; 16 | #endif 17 | -------------------------------------------------------------------------------- /tools/driver/include/linux/google/asic_fw_device_owner_accessor.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #ifndef _LINUX_GOOGLE_ASIC_FW_DEVICE_OWNER_ACCESSOR_H_ 6 | #define _LINUX_GOOGLE_ASIC_FW_DEVICE_OWNER_ACCESSOR_H_ 7 | #include 8 | inline unsigned long asic_fw_device_owner_value(const unsigned long *reg_value) 9 | { 10 | return (unsigned long)(((((*reg_value) >> 0) & 0xffffffffffffffffULL) 11 | << 0)); 12 | } 13 | inline int set_asic_fw_device_owner_value(unsigned long *reg_value, 14 | unsigned long value) 15 | { 16 | if (value & ~(0xffffffffffffffffULL)) 17 | return 1; 18 | (*reg_value) = ((*reg_value) & ~((0xffffffffffffffffULL) << 0)) | 19 | (((value >> 0) & (0xffffffffffffffffULL)) << 0); 20 | return 0; 21 | } 22 | #endif 23 | -------------------------------------------------------------------------------- /tools/driver/include/trace/events/gasket_interrupt.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #undef TRACE_SYSTEM 6 | #define TRACE_SYSTEM gasket_interrupt 7 | #if !defined(_TRACE_GASKET_INTERRUPT_H) || defined(TRACE_HEADER_MULTI_READ) 8 | #define _TRACE_GASKET_INTERRUPT_H 9 | #include 10 | TRACE_EVENT(gasket_interrupt_event, 11 | TP_PROTO(const char *name, 12 | int interrupt), 13 | TP_ARGS(name, interrupt), 14 | TP_STRUCT__entry( 15 | __string(device_name, name) 16 | __field(int, interrupt) 17 | ), 18 | TP_fast_assign( 19 | __assign_str(device_name, name); 20 | __entry->interrupt = interrupt; 21 | ), 22 | TP_printk("device %s, interrupt %d", 23 | __get_str(device_name), __entry->interrupt) 24 | ); 25 | #endif 26 | #include 27 | -------------------------------------------------------------------------------- /tools/driver/include/trace/events/gasket_mmap.h: -------------------------------------------------------------------------------- 1 | /* SPDX-License-Identifier: GPL-2.0 */ 2 | /* 3 | * Copyright (C) 2021 Google LLC. 4 | */ 5 | #undef TRACE_SYSTEM 6 | #define TRACE_SYSTEM gasket_mmap 7 | #if !defined(_TRACE_GASKET_MMAP_H) || defined(TRACE_HEADER_MULTI_READ) 8 | #define _TRACE_GASKET_MMAP_H 9 | #include 10 | TRACE_EVENT(gasket_mmap_entry, 11 | TP_PROTO(const char *name, 12 | ulong requested_offset, 13 | ulong requested_length), 14 | TP_ARGS(name, requested_offset, requested_length), 15 | TP_STRUCT__entry( 16 | __string(device_name, name) 17 | __field(ulong, requested_offset) 18 | __field(ulong, requested_length) 19 | ), 20 | TP_fast_assign( 21 | __assign_str(device_name, name); 22 | __entry->requested_offset = requested_offset; 23 | __entry->requested_length = requested_length; 24 | ), 25 | TP_printk("device %s, requested_offset 0x%lx, requested_length 0x%lx", 26 | __get_str(device_name), __entry->requested_offset, 27 | __entry->requested_length)); 28 | TRACE_EVENT(gasket_mmap_exit, 29 | TP_PROTO(uint retval), 30 | TP_ARGS(retval), 31 | TP_STRUCT__entry( 32 | __field(int, retval) 33 | ), 34 | TP_fast_assign( 35 | __entry->retval = retval; 36 | ), 37 | TP_printk("return value %d", __entry->retval)); 38 | #endif 39 | #include 40 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/examples/v4/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "region" {} 4 | variable "tpu_node_pools" {} 5 | variable "maintenance_interval" {} 6 | 7 | 8 | module "tpu-gke" { 9 | source = "../../module" 10 | project_id = var.project_id 11 | resource_name_prefix = var.resource_name_prefix 12 | region = var.region 13 | tpu_node_pools = var.tpu_node_pools 14 | maintenance_interval = var.maintenance_interval 15 | } 16 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/examples/v4/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "kubernetes_cluster_host" { 17 | value = module.tpu-gke.kubernetes_cluster_host 18 | description = "GKE Cluster Host" 19 | } 20 | 21 | output "nodepool_tpu_topology" { 22 | value = module.tpu-gke.nodepool_tpu_topology 23 | description = "GKE TPU topology" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/examples/v4/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | tpu_node_pools = [{ 5 | zone = "us-central2-b" 6 | node_count = 2 7 | machine_type = "ct4p-hightpu-4t" 8 | topology = "2x2x2" 9 | }] 10 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/examples/v5e/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "node_pool_prefix" {} 4 | variable "region" {} 5 | variable "tpu_node_pools" {} 6 | variable "maintenance_interval" {} 7 | variable "is_tpu_node_private" {} 8 | 9 | 10 | module "tpu-gke" { 11 | source = "../../module" 12 | project_id = var.project_id 13 | resource_name_prefix = var.resource_name_prefix 14 | node_pool_prefix = var.node_pool_prefix 15 | region = var.region 16 | tpu_node_pools = var.tpu_node_pools 17 | maintenance_interval = var.maintenance_interval 18 | is_tpu_node_private = var.is_tpu_node_private 19 | } 20 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/examples/v5e/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "placement_policy_names" { 17 | value = module.tpu-gke.placement_policy_names 18 | description = "GKE TPU Placement Policy Names" 19 | } 20 | 21 | output "is_tpu_node_private" { 22 | value = var.is_tpu_node_private 23 | description = "whether we want to make TPU node private" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/examples/v5e/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-v5e-test" 3 | node_pool_prefix = "rp1" 4 | region = "us-east5" 5 | is_tpu_node_private = false 6 | tpu_node_pools = [{ 7 | zone = "us-east5-b" 8 | node_count = 32 9 | machine_type = "ct5lp-hightpu-4t" 10 | topology = "8x16" 11 | policy = "sb-compact-rp1" 12 | disk_type = "pd-balanced" 13 | disk_size_gb = 120 14 | }] 15 | maintenance_interval = "PERIODIC" 16 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/module/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = google_container_node_pool.multihost_tpu[0].cluster 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "placement_policy_names" { 17 | value = flatten([ 18 | google_container_node_pool.multihost_tpu[*].placement_policy[0].policy_name 19 | ]) 20 | description = "GKE TPU Placement Policy Names" 21 | } 22 | 23 | output "is_tpu_node_private" { 24 | value = var.is_tpu_node_private 25 | description = "whether we want to make TPU node private" 26 | } 27 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/add_node_pool/module/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | tpu_node_pools = [{ 5 | zone = "us-central2-b" 6 | node_count = 4 7 | machine_type = "ct4p-hightpu-4t" 8 | topology = "2x2x4" 9 | }, { 10 | zone = "us-central2-b" 11 | node_count = 4 12 | machine_type = "ct4p-hightpu-4t" 13 | topology = "2x2x4" 14 | }, { 15 | zone = "us-central2-b" 16 | node_count = 2 17 | machine_type = "ct4p-hightpu-4t" 18 | topology = "2x2x2" 19 | }] 20 | maintenance_interval = "AS_NEEDED" 21 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/examples/v4/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "region" {} 4 | variable "tpu_node_pools" {} 5 | variable "cpu_node_pool" {} 6 | variable "maintenance_interval" {} 7 | 8 | 9 | module "tpu-gke" { 10 | source = "../../module" 11 | project_id = var.project_id 12 | resource_name_prefix = var.resource_name_prefix 13 | region = var.region 14 | tpu_node_pools = var.tpu_node_pools 15 | cpu_node_pool = var.cpu_node_pool 16 | maintenance_interval = var.maintenance_interval 17 | } 18 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/examples/v4/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "kubernetes_cluster_host" { 17 | value = module.tpu-gke.kubernetes_cluster_host 18 | description = "GKE Cluster Host" 19 | } 20 | 21 | output "nodepool_tpu_topology" { 22 | value = module.tpu-gke.nodepool_tpu_topology 23 | description = "GKE TPU topology" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/examples/v4/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | tpu_node_pools = [{ 5 | zone = "us-central2-b" 6 | node_count = 2 7 | machine_type = "ct4p-hightpu-4t" 8 | topology = "2x2x2" 9 | }] 10 | cpu_node_pool = { 11 | zone = ["us-central2-a", "us-central2-b", "us-central2-c"] 12 | machine_type = "n2-standard-8", 13 | initial_node_count_per_zone = 1, 14 | min_node_count_per_zone = 1, 15 | max_node_count_per_zone = 30, 16 | } 17 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/examples/v5e/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "region" {} 4 | variable "cpu_node_pool" {} 5 | variable "authorized_cidr_blocks" {} 6 | variable "is_cpu_node_private" {} 7 | 8 | 9 | module "tpu-gke" { 10 | source = "../../module" 11 | project_id = var.project_id 12 | resource_name_prefix = var.resource_name_prefix 13 | region = var.region 14 | cpu_node_pool = var.cpu_node_pool 15 | is_cpu_node_private = var.is_cpu_node_private 16 | authorized_cidr_blocks = var.authorized_cidr_blocks 17 | } 18 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/examples/v5e/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "authorized_cidr_blocks" { 17 | value = var.authorized_cidr_blocks 18 | description = "Cluster allowed cidr blocks " 19 | } 20 | 21 | output "is_cpu_node_private" { 22 | value = var.is_cpu_node_private 23 | description = "whether we want to make CPU node private" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/examples/v5e/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-v5e-test" 3 | region = "us-east5" 4 | authorized_cidr_blocks = [] 5 | is_cpu_node_private = false 6 | cpu_node_pool = { 7 | zone = ["us-east5-a", "us-east5-b", "us-east5-c"] 8 | machine_type = "n2-standard-8", 9 | initial_node_count_per_zone = 1, 10 | min_node_count_per_zone = 1, 11 | max_node_count_per_zone = 30, 12 | } 13 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/module/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = google_container_cluster.tpu_cluster.name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "authorized_cidr_blocks" { 17 | value = var.authorized_cidr_blocks 18 | description = "Cluster allowed cidr blocks " 19 | } 20 | 21 | output "is_cpu_node_private" { 22 | value = var.is_cpu_node_private 23 | description = "whether we want to make CPU node private" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_with_compact_placement/create_cluster/module/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | authorized_cidr_blocks = [] 5 | cpu_node_pool = { 6 | zone = ["us-central2-a", "us-central2-b", "us-central2-c"] 7 | machine_type = "n2-standard-64", 8 | initial_node_count_per_zone = 1, 9 | min_node_count_per_zone = 1, 10 | max_node_count_per_zone = 10 11 | } 12 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/examples/v4/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "region" {} 4 | variable "tpu_node_pools" {} 5 | variable "maintenance_interval" {} 6 | 7 | 8 | module "tpu-gke" { 9 | source = "../../module" 10 | project_id = var.project_id 11 | resource_name_prefix = var.resource_name_prefix 12 | region = var.region 13 | tpu_node_pools = var.tpu_node_pools 14 | maintenance_interval = var.maintenance_interval 15 | } 16 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/examples/v4/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "kubernetes_cluster_host" { 17 | value = module.tpu-gke.kubernetes_cluster_host 18 | description = "GKE Cluster Host" 19 | } 20 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/examples/v4/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | tpu_node_pools = [{ 5 | zone = "us-central2-b" 6 | node_count = 2 7 | machine_type = "ct4p-hightpu-4t" 8 | topology = "2x2x2" 9 | }] 10 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/examples/v5e/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "node_pool_prefix" {} 4 | variable "region" {} 5 | variable "tpu_node_pools" {} 6 | variable "maintenance_interval" {} 7 | variable "is_tpu_node_private" {} 8 | 9 | 10 | module "tpu-gke" { 11 | source = "../../module" 12 | project_id = var.project_id 13 | resource_name_prefix = var.resource_name_prefix 14 | node_pool_prefix = var.node_pool_prefix 15 | region = var.region 16 | tpu_node_pools = var.tpu_node_pools 17 | maintenance_interval = var.maintenance_interval 18 | is_tpu_node_private = var.is_tpu_node_private 19 | } 20 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/examples/v5e/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "is_tpu_node_private" { 17 | value = var.is_tpu_node_private 18 | description = "whether we want to make TPU node private" 19 | } 20 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/examples/v5e/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-v5e-test" 3 | node_pool_prefix = "batch1" 4 | region = "us-east5" 5 | is_tpu_node_private = false 6 | tpu_node_pools = [{ 7 | zone = "us-east5-b" 8 | node_count = 16 9 | machine_type = "ct5lp-hightpu-4t" 10 | topology = "8x8" 11 | disk_type = "pd-balanced" 12 | disk_size_gb = 120 13 | }] 14 | maintenance_interval = "PERIODIC" 15 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/module/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = google_container_node_pool.multihost_tpu[0].cluster 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "is_tpu_node_private" { 17 | value = var.is_tpu_node_private 18 | description = "whether we want to make TPU node private" 19 | } 20 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/add_node_pool/module/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | tpu_node_pools = [{ 5 | zone = "us-central2-b" 6 | node_count = 4 7 | machine_type = "ct4p-hightpu-4t" 8 | topology = "2x2x4" 9 | }, { 10 | zone = "us-central2-b" 11 | node_count = 4 12 | machine_type = "ct4p-hightpu-4t" 13 | topology = "2x2x4" 14 | }, { 15 | zone = "us-central2-b" 16 | node_count = 2 17 | machine_type = "ct4p-hightpu-4t" 18 | topology = "2x2x2" 19 | }] 20 | maintenance_interval = "AS_NEEDED" 21 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/examples/v4/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "region" {} 4 | variable "tpu_node_pools" {} 5 | variable "cpu_node_pool" {} 6 | variable "maintenance_interval" {} 7 | 8 | 9 | module "tpu-gke" { 10 | source = "../../module" 11 | project_id = var.project_id 12 | resource_name_prefix = var.resource_name_prefix 13 | region = var.region 14 | tpu_node_pools = var.tpu_node_pools 15 | cpu_node_pool = var.cpu_node_pool 16 | maintenance_interval = var.maintenance_interval 17 | } 18 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/examples/v4/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "kubernetes_cluster_host" { 17 | value = module.tpu-gke.kubernetes_cluster_host 18 | description = "GKE Cluster Host" 19 | } 20 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/examples/v4/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | tpu_node_pools = [{ 5 | zone = "us-central2-b" 6 | node_count = 2 7 | machine_type = "ct4p-hightpu-4t" 8 | topology = "2x2x2" 9 | }] 10 | cpu_node_pool = { 11 | zone = ["us-central2-a", "us-central2-b", "us-central2-c"] 12 | machine_type = "n2-standard-8", 13 | initial_node_count_per_zone = 1, 14 | min_node_count_per_zone = 1, 15 | max_node_count_per_zone = 30, 16 | } 17 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/examples/v5e/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "region" {} 4 | variable "cpu_node_pool" {} 5 | variable "authorized_cidr_blocks" {} 6 | variable "is_cpu_node_private" {} 7 | 8 | 9 | module "tpu-gke" { 10 | source = "../../module" 11 | project_id = var.project_id 12 | resource_name_prefix = var.resource_name_prefix 13 | region = var.region 14 | cpu_node_pool = var.cpu_node_pool 15 | is_cpu_node_private = var.is_cpu_node_private 16 | authorized_cidr_blocks = var.authorized_cidr_blocks 17 | } 18 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/examples/v5e/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "authorized_cidr_blocks" { 17 | value = var.authorized_cidr_blocks 18 | description = "Cluster allowed cidr blocks " 19 | } 20 | 21 | output "is_cpu_node_private" { 22 | value = var.is_cpu_node_private 23 | description = "whether we want to make CPU node private" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/examples/v5e/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-v5e-test" 3 | region = "us-east5" 4 | authorized_cidr_blocks = [] 5 | is_cpu_node_private = false 6 | cpu_node_pool = { 7 | zone = ["us-east5-a", "us-east5-b", "us-east5-c"] 8 | machine_type = "n2-standard-8", 9 | initial_node_count_per_zone = 1, 10 | min_node_count_per_zone = 1, 11 | max_node_count_per_zone = 30, 12 | } 13 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/module/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = google_container_cluster.tpu_cluster.name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "authorized_cidr_blocks" { 17 | value = var.authorized_cidr_blocks 18 | description = "Cluster allowed cidr blocks " 19 | } 20 | 21 | output "is_cpu_node_private" { 22 | value = var.is_cpu_node_private 23 | description = "whether we want to make CPU node private" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/batching_without_compact_placement/create_cluster/module/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | authorized_cidr_blocks = [] 5 | cpu_node_pool = { 6 | zone = ["us-central2-a", "us-central2-b", "us-central2-c"] 7 | machine_type = "n2-standard-64", 8 | initial_node_count_per_zone = 1, 9 | min_node_count_per_zone = 1, 10 | max_node_count_per_zone = 10 11 | } 12 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/non_batching_with_compact_placement/examples/v4/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "region" {} 4 | variable "tpu_node_pools" {} 5 | variable "maintenance_interval" {} 6 | 7 | 8 | module "tpu-gke" { 9 | source = "../../module" 10 | project_id = var.project_id 11 | resource_name_prefix = var.resource_name_prefix 12 | region = var.region 13 | tpu_node_pools = var.tpu_node_pools 14 | maintenance_interval = var.maintenance_interval 15 | } 16 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/non_batching_with_compact_placement/examples/v4/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "kubernetes_cluster_host" { 17 | value = module.tpu-gke.kubernetes_cluster_host 18 | description = "GKE Cluster Host" 19 | } 20 | 21 | output "nodepool_tpu_topology" { 22 | value = module.tpu-gke.nodepool_tpu_topology 23 | description = "GKE TPU topology" 24 | } 25 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/non_batching_with_compact_placement/examples/v4/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | tpu_node_pools = [{ 5 | zone = "us-central2-b" 6 | node_count = 2 7 | machine_type = "ct4p-hightpu-4t" 8 | topology = "2x2x2" 9 | }] 10 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/non_batching_with_compact_placement/examples/v5e/main.tf: -------------------------------------------------------------------------------- 1 | variable "project_id" {} 2 | variable "resource_name_prefix" {} 3 | variable "node_pool_prefix" {} 4 | variable "region" {} 5 | variable "cpu_node_pool" {} 6 | variable "tpu_node_pools" {} 7 | variable "maintenance_interval" {} 8 | variable "authorized_cidr_blocks" {} 9 | variable "is_cpu_node_private" {} 10 | variable "is_tpu_node_private" {} 11 | 12 | 13 | module "tpu-gke" { 14 | source = "../../module" 15 | project_id = var.project_id 16 | resource_name_prefix = var.resource_name_prefix 17 | node_pool_prefix = var.node_pool_prefix 18 | region = var.region 19 | cpu_node_pool = var.cpu_node_pool 20 | tpu_node_pools = var.tpu_node_pools 21 | maintenance_interval = var.maintenance_interval 22 | authorized_cidr_blocks = var.authorized_cidr_blocks 23 | is_cpu_node_private = var.is_cpu_node_private 24 | is_tpu_node_private = var.is_tpu_node_private 25 | } 26 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/non_batching_with_compact_placement/examples/v5e/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = module.tpu-gke.kubernetes_cluster_name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "kubernetes_cluster_host" { 17 | value = module.tpu-gke.kubernetes_cluster_host 18 | description = "GKE Cluster Host" 19 | } 20 | 21 | output "placement_policy_names" { 22 | value = module.tpu-gke.placement_policy_names 23 | description = "GKE TPU Placement Policy Names" 24 | } 25 | 26 | output "authorized_cidr_blocks" { 27 | value = var.authorized_cidr_blocks 28 | description = "Cluster allowed cidr blocks " 29 | } 30 | 31 | output "is_cpu_node_private" { 32 | value = var.is_cpu_node_private 33 | description = "whether we want to make CPU node private" 34 | } 35 | 36 | output "is_tpu_node_private" { 37 | value = var.is_tpu_node_private 38 | description = "whether we want to make TPU node private" 39 | } 40 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/non_batching_with_compact_placement/module/outputs.tf: -------------------------------------------------------------------------------- 1 | output "region" { 2 | value = var.region 3 | description = "GCloud Region" 4 | } 5 | 6 | output "project_id" { 7 | value = var.project_id 8 | description = "GCloud Project ID" 9 | } 10 | 11 | output "kubernetes_cluster_name" { 12 | value = google_container_cluster.tpu_cluster.name 13 | description = "GKE Cluster Name" 14 | } 15 | 16 | output "kubernetes_cluster_host" { 17 | value = google_container_cluster.tpu_cluster.endpoint 18 | description = "GKE Cluster Host" 19 | } 20 | 21 | output "placement_policy_names" { 22 | value = flatten([ 23 | google_container_node_pool.multihost_tpu[*].placement_policy[0].policy_name 24 | ]) 25 | description = "GKE TPU Placement Policy Names" 26 | } 27 | 28 | output "authorized_cidr_blocks" { 29 | value = var.authorized_cidr_blocks 30 | description = "Cluster allowed cidr blocks " 31 | } 32 | 33 | output "is_cpu_node_private" { 34 | value = var.is_cpu_node_private 35 | description = "whether we want to make CPU node private" 36 | } 37 | 38 | output "is_tpu_node_private" { 39 | value = var.is_tpu_node_private 40 | description = "whether we want to make TPU node private" 41 | } 42 | -------------------------------------------------------------------------------- /tools/kubernetes/terraform/non_batching_with_compact_placement/module/terraform.tfvars: -------------------------------------------------------------------------------- 1 | project_id = "project-id" 2 | resource_name_prefix = "tpu-test" 3 | region = "us-central2" 4 | authorized_cidr_blocks = [] 5 | cpu_node_pool = { 6 | zone = ["us-central2-a", "us-central2-b", "us-central2-c"] 7 | machine_type = "n2-standard-64", 8 | initial_node_count_per_zone = 1, 9 | min_node_count_per_zone = 1, 10 | max_node_count_per_zone = 10 11 | } 12 | tpu_node_pools = [{ 13 | zone = "us-central2-b" 14 | node_count = 4 15 | machine_type = "ct4p-hightpu-4t" 16 | topology = "2x2x4" 17 | }, { 18 | zone = "us-central2-b" 19 | node_count = 4 20 | machine_type = "ct4p-hightpu-4t" 21 | topology = "2x2x4" 22 | }, { 23 | zone = "us-central2-b" 24 | node_count = 2 25 | machine_type = "ct4p-hightpu-4t" 26 | topology = "2x2x2" 27 | }] 28 | maintenance_interval = "AS_NEEDED" 29 | -------------------------------------------------------------------------------- /tools/ray_tpu/deploy.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright 2023 The TensorFlow Authors. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2. 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================== 16 | 17 | RESOURCE_NAME=$USER-dev 18 | 19 | gcloud compute scp --recurse src/* ${RESOURCE_NAME}:~/src/ -------------------------------------------------------------------------------- /tools/ray_tpu/legacy/requirements-notebook.txt: -------------------------------------------------------------------------------- 1 | jinja2==3.0.3 2 | zipp==3.1.0 3 | ipyparallel==8.6.1 4 | jupyterlab==4.0.0 -------------------------------------------------------------------------------- /tools/ray_tpu/legacy/tpu_ray_cluster.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/tpu/c3186a4386eb090f8f13bb07cd4bae0b149b4e01/tools/ray_tpu/legacy/tpu_ray_cluster.png -------------------------------------------------------------------------------- /tools/ray_tpu/src/requirements.txt: -------------------------------------------------------------------------------- 1 | google-api-python-client 2 | google-auth-httplib2 3 | google-auth-oauthlib 4 | absl-py 5 | fabric==2.7.1 6 | patchwork 7 | ray[default] 8 | pyOpenSSL==23.0.0 -------------------------------------------------------------------------------- /tools/ray_tpu/src/serve/requirements.txt: -------------------------------------------------------------------------------- 1 | fastapi 2 | ray[serve]==2.5.1 3 | pillow -------------------------------------------------------------------------------- /tools/ray_tpu/src/serve/set_ray_address.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright 2023 The TensorFlow Authors. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================== 16 | 17 | echo "Make sure that you are running this as source ./set_ray_address.sh" 18 | EXTERNAL_HEAD_IP=$(ray get-head-ip cluster.yaml) 19 | 20 | # Getting the required IP using gcloud command and grep 21 | RAY_HEAD_IP=$(gcloud compute instances list | grep $EXTERNAL_HEAD_IP | awk '{print $4}') 22 | echo "Set RAY_HEAD_IP=${RAY_HEAD_IP}" 23 | 24 | # Setting the RAY_ADDRESS 25 | RAY_ADDRESS="http://${RAY_HEAD_IP}:8265" 26 | echo "Set RAY_ADDRESS=$RAY_ADDRESS" -------------------------------------------------------------------------------- /tools/ray_tpu/src/tune/requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py==1.0.0 2 | clu==0.0.6 3 | flax==0.4.1 4 | jax==0.3.4 5 | --find-links https://storage.googleapis.com/jax-releases/jax_releases.html 6 | jaxlib==0.3.2 7 | ml-collections==0.1.0 8 | numpy==1.22.0 9 | optax==0.1.0 10 | pandas 11 | ray[tune]==2.5.1 12 | tensorflow-cpu==2.9.3 13 | tensorflow-datasets==4.4.0 14 | protobuf==3.19.0 -------------------------------------------------------------------------------- /tools/ray_tpu/src/tune/set_ray_address.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright 2023 The TensorFlow Authors. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================== 16 | 17 | echo "Make sure that you are running this as source ./set_ray_address.sh" 18 | EXTERNAL_HEAD_IP=$(ray get-head-ip cluster.yaml) 19 | 20 | # Getting the required IP using gcloud command and grep 21 | RAY_HEAD_IP=$(gcloud compute instances list | grep $EXTERNAL_HEAD_IP | awk '{print $4}') 22 | echo "Set RAY_HEAD_IP=${RAY_HEAD_IP}" 23 | 24 | # Setting the RAY_ADDRESS 25 | RAY_ADDRESS="http://${RAY_HEAD_IP}:8265" 26 | echo "Set RAY_ADDRESS=$RAY_ADDRESS" -------------------------------------------------------------------------------- /tools/retry/README.md: -------------------------------------------------------------------------------- 1 | # TPU Retry script 2 | 3 | ## About 4 | This is a toy example of a script that can be used to poll for TPU health and delete/create if stuck in an unhealthy state. 5 | 6 | Please continue reading for expectations: 7 | - Cloud TPUs are expected to undergo maintenance events, but is expected to recover. 8 | - Generally, saving checkpoints more often (at least every hour) allows you to gracefully recover. Also ensure that your training script resumes from checkpoint correctly. 9 | - In unexpected circumstances, it's possible that the TPU does not recover from maintenance event. This script showcases an example of how to detect this and delete/re-create the TPU. 10 | - If your process running the training script crashes, you can modify this script to re-try running the train script. 11 | 12 | 13 | ## Example usage 14 | You can run this script on your VM. This assumes that you have already exported the tpu name, e.g. 15 | ``` 16 | export TPU_NAME={my_tpu} 17 | ``` 18 | 19 | Within your VM you can run this script: 20 | ``` 21 | ./retry.sh $TPU_NAME & 22 | ``` 23 | 24 | You are free to modify the polling frequency, the re-creation logic, etc. 25 | 26 | Note that this runs indefinitely, so you will need to `pkill` the script once you are done. 27 | 28 | *NOTE*: Please modify the script to reflect the correct TPU deletion/creation command, as this will differ if you're using e.g. TPU VM or a reserved TPU. 29 | --------------------------------------------------------------------------------