Google 的 Cheng Xing 和 Michael Broughton 发表
训练大型机器学习模型是 TensorFlow 的核心能力。多年来，规模已成为 NLP、图像识别、药物发现等许多现代机器学习系统的重要特征。利用多台机器来提高计算能力和吞吐量，已在该领域取得了巨大进步。同样在量子计算和量子机器学习中，更多机器资源的可用性加速了对更大量子态和更复杂系统的模拟。在本教程中，您将学习如何使用 TensorFlow 和 TensorFlow Quantum 来进行大规模分布式 QML 模拟。运行具有更大 FLOP/s 计数的更大模拟，为以前在更小规模下无法实现的研究打开了新的可能性。在下图中，我们概述了用于量子模拟的几种不同硬件设置的大致扩展能力。
从我们在多工作器环境中的实验来看，训练一个包含 1,000 个训练示例的 23 量子比特 QCNN（对应于使用全状态向量模拟模拟的大约 3,000 个电路）在 32 节点 (512 个 vCPU) 集群上每个 epoch 需要 5 分钟，花费几美元。相比之下，在单个工作器上进行相同的训练工作每个 epoch 大约需要 4 个小时。更进一步，使用超过 10,000 个虚拟 CPU，可以在几个小时内运行数十万个 30 量子比特电路，这在单工作器环境中可能需要几周才能完成。实际性能和成本可能因您的云设置而异，例如 VM 机器类型、集群总运行时间等。在执行更大规模的实验之前，我们建议先从一个小集群开始，比如本教程中使用的一个。
本教程的源代码可在 TensorFlow Quantum GitHub 存储库中找到。README.md 包含使本教程快速运行的最快方法。本教程将详细介绍每个步骤，以帮助您理解基础概念并将它们集成到自己的项目中。让我们开始吧！
第一步是在 Google Cloud 中创建基础设施资源。如果您已有现有的 Google Cloud 环境，则具体步骤可能会有所不同，例如由于组织策略限制。这是一套最常见的必要步骤的指南。请注意，您将为创建的 Google Cloud 资源付费，这里 是本教程中使用的可计费资源的摘要。如果您是 Google Cloud 新用户，您有资格获得 $300 的信用额度。如果您是学术机构的一部分，您可能有资格获得 Google Cloud 研究信用额度。
在本教程中，您将运行多个 shell 命令。为此，您可以使用计算机上可用的本地 Unix shell 或 Cloud Shell，Cloud Shell 中已经包含后面提到的许多工具。
一个自动执行以下步骤的脚本可在 setup.sh 中找到。本节将详细介绍每个步骤，如果您是第一次使用 Google Cloud，我们建议您完整地学习本节。如果您希望自动执行 Google Cloud 设置过程并跳过本节
setup.sh 并配置内部的参数值。./setup.sh infra。在本教程中，您将使用一些 Google Cloud 产品
要准备好您的云环境，请先按照以下快速入门指南操作
在本教程中，您可以在创建集群的说明之前停止 Kubernetes Engine 快速入门。此外，请安装 gsutil（Cloud Storage 命令行工具）（如果您使用的是 Cloud Shell，则 gsutil 已安装）
gcloud components install gsutil作为参考，本教程中的 shell 命令将引用以下变量。其中一些变量在教程后面在每个命令的上下文中会更有意义。
${CLUSTER_NAME}：您在 Google Kubernetes Engine 上首选的 Kubernetes 集群名称。${PROJECT}：您的 Google Cloud 项目 ID。${NUM_NODES}：集群中的 VM 数量。${MACHINE_TYPE}：VM 的 机器类型。这控制着每个 VM 的 CPU 和内存资源量。${SERVICE_ACCOUNT_NAME}：Google Cloud IAM 服务帐户和关联的 Kubernetes 服务帐户的名称。${ZONE}：Kubernetes 集群的 Google Cloud 区域。${BUCKET_REGION}：Google Cloud 存储桶的 Google Cloud 区域。${BUCKET_NAME}：用于存储训练输出的 Google Cloud 存储桶的名称。owner，或者具有以下所有角色container.admin
iam.serviceAccountAdmin
storage.admin
要检查您的角色，请运行
gcloud projects get-iam-policy ${PROJECT}完成快速入门指南后，运行以下命令创建 Kubernetes 集群
gcloud container clusters create ${CLUSTER_NAME} --workload-pool=${PROJECT}.svc.id.goog --num-nodes=${NUM_NODES} --machine-type=${MACHINE_TYPE} --zone=${ZONE} --preemptible
--num-nodes 是支持 Kubernetes 集群的 Compute Engine 虚拟机数量。这并不一定与您希望用于 QCNN 作业的工作器副本数量相同，因为 Kubernetes 能够在同一节点上调度多个副本，前提是节点有足够的 CPU 和内存资源。如果您是第一次尝试本教程，我们建议使用 2 个节点。
--machine-type 指定 VM 的 机器类型。如果您是第一次尝试本教程，我们建议使用“n1-standard-2”，它具有 2 个 vCPU 和 7.5 GB 的内存。
--zone 是您想要运行集群的 Google Cloud 区域（例如“us-west1-a”）。
--workload-pool 启用 GKE Workload Identity 功能，该功能将 Kubernetes 服务帐户 与 Google Cloud IAM 服务帐户 关联起来。为了实现细粒度的访问控制，建议使用 IAM 服务帐户来访问各种 Google Cloud 产品。在这里，您将创建一个服务帐户，以便由您的 QCNN 作业使用。Kubernetes 服务帐户是将此 IAM 服务帐户的凭据注入到您的工作器容器中的机制。
--preemptible 使用 Compute Engine 可抢占式 VM 来支持 Kubernetes 集群。与常规 VM 相比，它们的成本低至 80%，但缺点是 VM 可能会在任何时候被抢占，这将终止训练过程。这非常适合使用大型集群进行短时运行的训练会话。
然后，您可以创建一个 IAM 服务帐户
gcloud iam service-accounts create ${SERVICE_ACCOUNT_NAME}gcloud iam service-accounts add-iam-policy-binding --role roles/iam.workloadIdentityUser --member "serviceAccount:${PROJECT}.svc.id.goog[default/${SERVICE_ACCOUNT_NAME}]" ${SERVICE_ACCOUNT_NAME}@${PROJECT}.iam.gserviceaccount.comgsutil mb -p ${PROJECT} -l ${BUCKET_REGION} -b on gs://${BUCKET_NAME}为了使您的 Cloud Storage 数据可供 QCNN 作业访问，请向您的 IAM 服务帐户授予权限
gsutil iam ch serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT}.iam.gserviceaccount.com:roles/storage.admin gs://${BUCKET_NAME}云环境搭建完成后，您现在可以将必要的 Kubernetes 工具安装到集群中。您需要 `tf-operator`，它是 KubeFlow 的一个组件。KubeFlow 是一个用于在 Kubernetes 上运行机器学习工作负载的工具包，而 `tf-operator` 是一个子组件，它简化了 TensorFlow 作业的管理。`tf-operator` 可以单独安装，无需安装更大的 KubeFlow 安装程序。
要安装 `tf-operator`，请运行
docker pull k8s.gcr.io/kustomize/kustomize:v3.10.0
docker run k8s.gcr.io/kustomize/kustomize:v3.10.0 build "github.com/kubeflow/tf-operator.git/manifests/overlays/standalone?ref=v1.1.0" | kubectl apply -f -您现在可以获取 TensorFlow Quantum 研究分支上的 QCNN 代码 并将其准备以分布式方式运行。让我们克隆源代码
git clone https://github.com/tensorflow/quantum.git && cd quantum && git checkout origin/research && cd qcnn_multiworkergit clone [email protected]:tensorflow/quantum.git && cd quantum && git checkout origin/research && cd qcnn_multiworker`training` 目录包含执行 QCNN 分布式训练所需的组件。`training/qcnn.py` 和 `common/qcnn_common.py` 的组合与 TensorFlow Quantum 中的混合 QCNN 示例相同，但添加了一些功能
MultiWorkerMirroredStrategy 是 TensorFlow 中执行同步分布式训练的机制。您的现有模型已通过几行额外的代码增强，以支持分布式训练。
在 `training/qcnn.py` 的开头，我们设置了 MultiWorkerMirroredStrategy
strategy = tf.distribute.MultiWorkerMirroredStrategy()... = qcnn_common.prepare_model(strategy)QCNN 分布式训练作业的每个工作器将运行此 Python 代码的副本。每个工作器都需要知道所有其他工作器的网络端点。通常使用 `TF_CONFIG` 环境变量来实现此目的，但在我们的示例中，`tf-operator` 会在后台自动注入它。
模型训练完成后，权重将上传到您的 Cloud Storage 存储桶，以便稍后由推理作业访问。
if task_type == 'worker' and task_id == 0:
qcnn_weights_path='/tmp/qcnn_weights.h5'
qcnn_model.save_weights(qcnn_weights_path)
upload_blob(args.weights_gcs_bucket, qcnn_weights_path, f'qcnn_weights.h5')在继续 Kubernetes 部署设置并启动您的工作器之前，需要在教程源代码中配置几个参数以匹配您的设置。提供的脚本 `setup.sh` 可用于简化此过程。
打开 `setup.sh` 并配置其中的参数值，如果您在之前的步骤中尚未完成此操作。然后运行
./setup.sh parammake training本节的其余部分将详细介绍 Kubernetes 设置。
在 Kubernetes 中以 容器 运行之前，QCNN 作业需要使用 Docker 打包成容器镜像并上传到容器注册表。`Dockerfile` 包含镜像规范。要构建和上传镜像，请运行
docker build -t gcr.io/${PROJECT}/qcnn .
docker push gcr.io/${PROJECT}/qcnnapiVersion: v1
kind: ServiceAccount
metadata:
annotations:
iam.gke.io/gcp-service-account: ${SERVICE_ACCOUNT_NAME}@${PROJECT}.iam.gserviceaccount.com
name: ${SERVICE_ACCOUNT_NAME}kubectl apply -f common/sa.yaml最后一步是创建分布式训练作业。`training/qcnn.yaml` 包含作业的 Kubernetes 规范。在 Kubernetes 中，具有相关功能的多个容器被分组到一个称为 Pod 的单个实体中，它是可以安排的最基本的作业单元。通常，用户利用现有的资源类型（如 Deployment 和 Job）来创建和管理工作负载。您将使用 `TFJob`（如 `kind` 字段中指定的那样），它不是 Kubernetes 内置的资源类型，而是 `tf-operator` 提供的 自定义资源，这使得更容易使用 TensorFlow 工作负载。
值得注意的是，`TFJob` 规范包含字段 `tfReplicaSpecs.Worker`，它允许您配置 MultiWorkerMirroredStrategy 工作器。`PS`（参数服务器）、`Chief` 和 `Evaluator` 的值也支持异步和其他形式的分布式训练。在幕后，`tf-operator` 为每个工作器副本创建两个 Kubernetes 资源
 |
| TFJob 为每个工作器副本生成一个服务和一个 Pod。一旦 TFJob 更新，更改将反映在底层服务和 Pod 中。工作器状态也会在 TFJob 中报告。 |
 |
| 服务将工作器服务器公开给集群中的其他部分。每个工作器通过使用目标工作器的服务名称作为 DNS 名称来与其他工作器通信。 |
在工作器规范中，有几个值得注意的字段
kubectl apply -f training/qcnn.yamlNAME READY STATUS RESTARTS
qcnn-worker-0 1/1 Running 0
qcnn-worker-1 1/1 Running 0kubectl logs <worker_pod_name>…
I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:411] Started server with target: grpc:/
/qcnn-worker-0.default.svc:2222
…
I tensorflow/core/profiler/rpc/profiler_server.cc:46] Profiler server listening on [::]:2223 selecte
d port:2223
…
Epoch 1/50
…
4/4 [==============================] - 7s 940ms/step - loss: 0.9387 - accuracy: 0.0000e+00 - val_loss: 0.7432 - val_accuracy: 0.0000e+00
…
I tensorflow/core/profiler/lib/profiler_session.cc:71] Profiler session collecting data.
I tensorflow/core/profiler/lib/profiler_session.cc:172] Profiler session tear down.
…
Epoch 50/50
4/4 [==============================] - 1s 222ms/step - loss: 0.1468 - accuracy: 0.4101 - val_loss: 0.2043 - val_accuracy: 0.4583
File /tmp/qcnn_weights.h5 uploaded to qcnn_weights.h5.`qcnn-worker-1` 的输出应该类似，只是最后一行缺失。首席工作器（工作器 0）负责保存整个模型的权重。
您还可以通过访问 Cloud Console 中的存储浏览器 并浏览您之前创建的存储桶来验证模型权重是否已保存。
要删除训练作业，请运行
kubectl delete -f training/qcnn.yamltensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=args.logdir,
histogram_freq=1,
update_freq=1,
profile_batch='10, 20')
…
history = qcnn_model.fit(x=train_excitations,
y=train_labels,
batch_size=32,
epochs=50,
verbose=1,
validation_data=(test_excitations, test_labels),
callbacks=[tensorboard_callback])tf.profiler.experimental.server.start(args.profiler_port)qcnn_model.compile(..., metrics=[‘accuracy’])
TensorFlow 分析器 是一个有助于调试模型训练作业性能瓶颈的工具。
在本教程中，我们使用编程模式（其中分析是在预定义的训练步骤范围内完成的）和采样模式（其中分析可以在按需完成）。对于 MultiWorkerMirroredStrategy 设置，目前编程模式仅输出首席（工作器 0）的分析数据，而采样模式能够分析所有工作器。
首次打开分析器时，显示的数据来自编程模式。概述页面可以让您了解每一步训练花费的时间。这将作为您尝试使用不同的方法来提高训练性能时的参考，无论是通过扩展基础设施（向集群添加更多 VM、使用具有更多 CPU 和内存的 VM、与硬件加速器集成）还是提高代码效率。
要查看 TensorBoard UI，您可以在 Kubernetes 中创建 TensorBoard 实例。Kubernetes 设置位于 `training/tensorboard.yaml` 中。此文件包含两个对象
也可以在您的工作站上运行本地 TensorBoard 实例，方法是将 `--logdir` 指向同一个 Cloud Storage 存储桶，尽管需要额外的 IAM 权限设置。
创建此 Kubernetes 设置
kubectl apply -f training/tensorboard.yaml在 `kubectl get pods` 的输出中，您应该看到有一个以 `qcnn-tensorboard` 为前缀的 Pod，它最终处于 Running 状态。要获取 TensorBoard 实例的 IP 地址，请运行
kubectl get svc tensorboard-service -wNAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S)
tensorboard-service LoadBalancer 10.123.240.9 <pending> 5001:32200/TCP负载均衡器需要一些时间才能配置，因此您可能不会立即看到 IP。一旦可用，请在浏览器中访问 `<ip>:5001` 以访问 TensorBoard UI。
在 TensorFlow 2.4 及更高版本中，可以在采样模式下分析多个工作器：可以通过在 Tensorboard 分析器中单击“Capture Profile”并将“Profile Service URL”设置为 `qcnn-worker-<replica_id>:2223` 来分析正在运行的训练作业中的工作器。要启用此功能，需要通过工作器服务公开分析器端口。教程源代码提供了一个脚本，用于使用分析器端口修补 `TFJob` 生成的所有工作器服务。运行
training/apply_profiler_ports.shinference/ 目录中。主文件 qcnn_inference.py 大部分重复使用 common/qcnn_common.py 中的模型构建代码，但会从您的 Cloud Storage 存储桶中加载模型权重。qcnn_weights_path = '/tmp/qcnn_weights.h5'
download_blob(args.weights_gcs_bucket, args.weights_gcs_path, qcnn_weights_path)
qcnn_model.load_weights(qcnn_weights_path)results = qcnn_model(test_excitations).numpy().flatten()
loss = tf.keras.losses.mean_squared_error(test_labels, results)make inference推理程序已构建到训练步骤中的 Docker 镜像中，因此您无需在此处构建新的镜像。推理作业规范 inference/inference.yaml 包含一个作业，其 Pod 规范指向该镜像，但会执行 qcnn_inference.py。运行 kubectl apply -f inference/inference.yaml 来创建作业。
以 inference-qcnn 为前缀的 Pod 最终应处于运行状态 (kubectl get pods)。在推理 Pod 的日志输出中 (kubectl logs <pod_name>)，均方误差应接近 TensorBoard UI 中显示的最终损失。
…
Blob qcnn_weights.h5 downloaded to /tmp/qcnn_weights.h5.
[-0.8220097 0.40201923 -0.82856977 0.46476707 -1.1281478 0.23317486
0.00584182 1.3351855 0.35139582 -0.09958048 1.2205497 -1.3038696
1.4065738 -1.1120421 -0.01021352 1.4553616 -0.70309246 -0.0518395
1.4699622 -1.3712595 -0.01870352 1.2939589 1.2865802 0.847203
0.3149605 1.1705848 -1.0051676 1.2537074 -0.2943283 -1.3489063
-1.4727883 1.4566276 1.3417912 0.9123422 0.2942805 -0.791862
1.2984066 -1.1139404 1.4648925 -1.6311806 -0.17530376 0.70148027
-1.0084027 0.09898916 0.4121615 0.62743163 -1.4237025 -0.6296255 ]
Test Labels
[-1 1 -1 1 -1 1 1 1 1 -1 1 -1 1 -1 1 1 -1 -1 1 -1 -1 1 1 1
1 1 -1 1 -1 -1 -1 1 1 1 -1 -1 1 -1 1 -1 -1 1 -1 1 1 1 -1 -1]
Mean squared error: tf.Tensor(0.29677835, shape=(), dtype=float32)这结束了我们对分布式训练之旅的介绍！在您完成本教程的实验后，本节将引导您完成清理 Google Cloud 资源的步骤。
首先，删除 Kubernetes 部署。运行
make delete-inference
kubectl delete -f training/tensorboard.yamlmake delete-traininggcloud container clusters delete ${CLUSTER_NAME} --zone=${ZONE}gsutil rm -r gs://${BUCKET_NAME}qcnn。现在您已经尝试了多工作程序设置，请尝试在您的项目中进行设置！由于本教程中提到的所有工具都在不断发展，因此使用多个工作程序进行训练的最佳实践会随着时间的推移而发生变化。请查看 TensorFlow Quantum GitHub 存储库中教程目录的更新！
随着您继续扩展实验，您最终可能会遇到基础设施限制，这些限制要求对本教程中使用的技术进行高级配置，因为在分布式环境中工作很复杂。要深入了解其中的一些，请查看以下资源
