GoogleCloudPlatform / GoogleCloudPlatform/pubsec-declarative-toolkit
GCP TPU (Tensor Processing Unit) project template for TensorFlow 2.x LLM workload training/inference enablement via the landing zone - specifically TPUv5
- Dominant language
- Shell
- Stars
- 36
- Forks
- 26
- PR merge metrics
- No merged PRs in 30d
Description
## Bootstrap TPU project
- follow https://cloud.google.com/tpu/docs/run-calculation-tensorflow
- https://github.com/ObrienlabsDev/machine-learning
```
FROM tensorflow/tensorflow:latest-gpu
WORKDIR /src
COPY /src/tflow.py .
CMD ["python", "tflow.py"]
```
- tflow.py
```
import tensorflow as tf
# modify for TF - TPUStrategy
# https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy#used-in-the-notebooks
strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
cifar = tf.keras.datasets.cifar100
(x_train, y_train), (x_test, y_test) = cifar.load_data()
with strategy.scope():
# https://www.tensorflow.org/api_docs/python/tf/keras/applications/resnet50/ResNet50
# https://keras.io/api/models/model/
parallel_model = tf.keras.applications.ResNet50(
include_top=True,
weights=None,
input_shape=(32, 32, 3),
classes=100,)
# https://saturncloud.io/blog/how-to-do-multigpu-training-with-keras/
#parallel_model = multi_gpu_model(model, gpus=2)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)
# https://keras.io/api/models/model_training_apis/
parallel_model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])
parallel_model.fit(x_train, y_train, epochs=10, batch_size=256)#5120)#7168)#7168)
```
## Reference
- detail comparison between NVidia tensor cores in the latest on-prem https://resources.nvidia.com/en-us-design-viz-stories-ep/rtx-5000-ada-datasheet?lx=CCKW39&contentType=data-sheet
- detail comparison between GCP L4 VMs with NVidia tensor cores https://cloud.google.com/blog/products/compute/introducing-g2-vms-with-nvidia-l4-gpus
Contributor guide
Assessment
This issue has not been assessed yet.