[Vertex AI SDK] Add custom serializer args for bigframes tensorflow (#2633)

* Add custom serializer args for bigframes tensorflow

* Add custom serializer args for remote prediction
This commit is contained in:
Matthew Tang
2024-01-24 20:47:02 +00:00
committed by GitHub
parent 556cbe7049
commit 0ea0c04195
@@ -456,8 +456,7 @@
" \"virginica\": 1,\n",
" \"setosa\": 2,\n",
"}\n",
"df[\"target\"] = df[\"species\"].map(species_categories)\n",
"df = df.drop(columns=[\"species\"])\n",
"df[\"species\"] = df[\"species\"].map(species_categories)\n",
"\n",
"train, test = bf_train_test_split(df, test_size=0.2)"
]
@@ -502,6 +501,12 @@
"model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_K80\"\n",
"model.fit.vertex.remote_config.accelerator_count = 4\n",
"\n",
"# (Optional) Set batch_size, target_col\n",
"model.fit.vertex.remote_config.serializer_args[train] = {\n",
" \"batch_size\": 32,\n",
" \"target_col\": \"species\",\n",
"}\n",
"\n",
"# Train model on Vertex\n",
"model.fit(train, epochs=10)"
]
@@ -531,6 +536,12 @@
"# Disable GPU for remote prediction\n",
"model.predict.vertex.remote_config.enable_cuda = False\n",
"\n",
"# (Optional) Set batch_size, target_col\n",
"model.predict.vertex.remote_config.serializer_args[train] = {\n",
" \"batch_size\": 32,\n",
" \"target_col\": \"species\",\n",
"}\n",
"\n",
"predictions = model.predict(train)\n",
"\n",
"print(f\"Remote predictions: {predictions}\")"
@@ -557,7 +568,7 @@
"source": [
"# User must convert bigframes to pandas dataframe for local evaluation\n",
"feature_columns = [\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]\n",
"label_columns = [\"target\"]\n",
"label_columns = [\"species\"]\n",
"\n",
"train_X_np = train[feature_columns].to_pandas().values.astype(float)\n",
"train_y_np = train[label_columns].to_pandas().values.astype(float)\n",