add more flexibility to training

cloud-annotations · Jan 5, 2020 · 2770f08 · 2770f08
1 parent 2fe1a13
commit 2770f08
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/cacli/cmd/train.go b/cacli/cmd/train.go
@@ -53,6 +53,10 @@ func init() {
 	trainCmd.Flags().String("name", "", "Optional project name")
 	trainCmd.Flags().String("output", "", "Optional output bucket")
 	trainCmd.Flags().Int("steps", 1000, "Number of training steps")
-	trainCmd.Flags().String("gpu", "k80", "k80 | k80x2 | k80x4 | v100 | v100x2")
+	trainCmd.Flags().String("gpu", "k80", "k80x2, k80x4, v100, v100x2")
 	trainCmd.Flags().String("script", "", "Custom training script.zip")
+
+	trainCmd.Flags().String("framework", "tensorflow", "keras, pytorch, caffe")
+	trainCmd.Flags().String("frameworkv", "1.14", "Framework version")
+	trainCmd.Flags().String("pythonv", "3.6", "Python version")
 }
diff --git a/cacli/cmd/train/train.go b/cacli/cmd/train/train.go
@@ -31,6 +31,10 @@ func Run(cmd *cobra.Command, args []string) {
 	gpu, err := cmd.Flags().GetString("gpu")
 	script, err := cmd.Flags().GetString("script")
 
+	framework, err := cmd.Flags().GetString("framework")
+	frameworkVersion, err := cmd.Flags().GetString("frameworkv")
+	pythonVersion, err := cmd.Flags().GetString("pythonv")
+
 	if err != nil {
 		e.Exit(err)
 	}
@@ -109,7 +113,7 @@ func Run(cmd *cobra.Command, args []string) {
 		// if non default steps, include it in project name.
 		projectName = projectName + " (" + strconv.Itoa(steps) + ")"
 	}
-	model, err := session.StartTraining(script, projectName, trainingBucket, outputBucket, steps, gpu)
+	model, err := session.StartTraining(script, projectName, trainingBucket, outputBucket, steps, gpu, framework, frameworkVersion, pythonVersion)
 	if err != nil {
 		e.Exit(err)
 	}

diff --git a/cacli/ibmcloud/wrapper.go b/cacli/ibmcloud/wrapper.go
@@ -332,17 +332,17 @@ func (s *AccountSession) CreateCredential(params CreateCredentialParams) (*Crede
 	return credential, nil
 }
 
-func (s *CredentialSession) StartTraining(trainingZip string, projectName string, bucket *s3.BucketExtended, output *s3.BucketExtended, steps int, gpu string) (*Model, error) {
+func (s *CredentialSession) StartTraining(trainingZip string, projectName string, bucket *s3.BucketExtended, output *s3.BucketExtended, steps int, gpu string, framework string, frameworkVersion string, pythonVersion string) (*Model, error) {
 	// TODO: We shouldn't hard code all of this.
 	trainingDefinition := &TrainingDefinition{
 		Name: projectName,
 		Framework: Framework{
-			Name:    "tensorflow",
-			Version: "1.12",
+			Name:    framework,
+			Version: frameworkVersion,
 			Runtimes: []Runtimes{
 				Runtimes{
 					Name:    "python",
-					Version: "3.6",
+					Version: pythonVersion,
 				},
 			},
 		},