Fine Tune Training (#1376)

* get models for web version

* fine-tune front-end and web command arguments

* updating user trained models

* update run trianing for tuneModel arguments

* electorn support

* pydantic type updates

* fix arguments for training

* support running locally on single gpu pc's by default

* update fine tune model web display

* Update client/platform/web-girder/views/RunTrainingMenu.vue

Co-authored-by: Michael Nagler <[email protected]>

* Update server/dive_server/crud_rpc.py

Co-authored-by: Michael Nagler <[email protected]>

* adressing comments

---------

Co-authored-by: Michael Nagler <[email protected]>

Author: BryonLewis

client/dive-common/apispec.ts

@@ -45,8 +45,16 @@ interface Category {
 }
 
 interface TrainingConfigs {
-  configs: string[];
-  default: string;
+  training: {
+    configs: string[];
+    default: string;
+  };
+  models: Record<string, {
+    name: string;
+    type: string;
+    path?: string;
+    folderId?: string;
+  }>;
 }
 
 type Pipelines = Record<string, Category>;
@@ -155,6 +163,12 @@ interface Api {
     config: string,
     annotatedFramesOnly: boolean,
     labelText?: string,
+    fineTuneModel?: {
+      name: string;
+      type: string;
+      path?: string;
+      folderId?: string;
+    },
   ): Promise<unknown>;
 
   loadMetadata(datasetId: string): Promise<DatasetMeta>;

client/platform/desktop/backend/native/common.ts

@@ -423,14 +423,36 @@ async function getPipelineList(settings: Settings): Promise<Pipelines> {
   return ret;
 }
 
+// Function to recursively traverse a directory and collect files with specified extensions
+function getFilesWithExtensions(dir: string, extensions: string[], fileList: string[] = []) {
+  const files = fs.readdirSync(dir);
+
+  files.forEach((file) => {
+    const filePath = npath.join(dir, file);
+    const fileStat = fs.statSync(filePath);
+
+    if (fileStat.isDirectory()) {
+      fileList.concat(getFilesWithExtensions(filePath, extensions, fileList));
+    } else {
+      const fileExtension = npath.extname(file).toLowerCase();
+      if (extensions.includes(fileExtension)) {
+        fileList.push(filePath);
+      }
+    }
+  });
+
+  return fileList;
+}
+
 /**
  * get training configurations
  */
 async function getTrainingConfigs(settings: Settings): Promise<TrainingConfigs> {
   const pipelinePath = npath.join(settings.viamePath, 'configs/pipelines');
-  const defaultTrainingConfiguration = 'train_detector_default.conf';
-  const allowedPatterns = /train_.*\.conf$/;
-  const disallowedPatterns = /.*(_nf|\.continue)\.viame_csv\.conf$|.*\.kw18\.conf$|.*\.habcam\.conf$|.*\.continue\.conf$/;
+  const defaultTrainingConfiguration = 'train_detector_default.viame_csv.conf';
+  const allowedPatterns = /\.viame_csv\.conf$/;
+  const disallowedPatterns = /.*(_nf|\.continue)\.viame_csv\.conf$/;
+  const allowedModelExtensions = ['.zip', '.pth', '.pt', '.py', '.weights', '.wt'];
   const exists = await fs.pathExists(pipelinePath);
   if (!exists) {
     throw new Error(`Path does not exist: ${pipelinePath}`);
@@ -439,9 +461,22 @@ async function getTrainingConfigs(settings: Settings): Promise<TrainingConfigs>
   configs = configs
     .filter((p) => (p.match(allowedPatterns) && !p.match(disallowedPatterns)))
     .sort((a, b) => (a === defaultTrainingConfiguration ? -1 : a.localeCompare(b)));
+  // Get Model files in the pipeline directory
+  const modelList = getFilesWithExtensions(pipelinePath, allowedModelExtensions);
+  const models: TrainingConfigs['models'] = {};
+  modelList.forEach((model) => {
+    models[npath.basename(model)] = {
+      name: npath.basename(model),
+      type: npath.extname(model),
+      path: model,
+    };
+  });
   return {
-    default: configs[0],
-    configs,
+    training: {
+      default: configs[0],
+      configs,
+    },
+    models,
   };
 }
 

client/platform/desktop/backend/native/viame.ts

@@ -394,6 +394,11 @@ async function train(
     command.push('--gt-frames-only');
   }
 
+  if (runTrainingArgs.fineTuneModel && runTrainingArgs.fineTuneModel.path) {
+    command.push('--init-weights');
+    command.push(runTrainingArgs.fineTuneModel.path);
+  }
+
   const job = observeChild(spawn(command.join(' '), {
     shell: viameConstants.shell,
     cwd: jobWorkDir,

client/platform/desktop/constants.ts

@@ -171,6 +171,13 @@ export interface RunTraining {
   annotatedFramesOnly: boolean;
   // contents of labels.txt file
   labelText?: string;
+  // fine tuning model
+  fineTuneModel?: {
+    name: string;
+    type: string;
+    path?: string;
+    folderId?: string;
+  };
 }
 
 export interface ConversionArgs {

client/platform/desktop/frontend/api.ts

@@ -110,13 +110,20 @@ async function runTraining(
   config: string,
   annotatedFramesOnly: boolean,
   labelText?: string,
+  fineTuneModel?: {
+    name: string;
+    type: string;
+    path?: string;
+    folderId?: string;
+  },
 ): Promise<DesktopJob> {
   const args: RunTraining = {
     datasetIds: folderIds,
     pipelineName,
     trainingConfig: config,
     annotatedFramesOnly,
     labelText,
+    fineTuneModel,
   };
   return ipcRenderer.invoke('run-training', args);
 }

client/platform/desktop/frontend/components/MultiTrainingMenu.vue

@@ -52,9 +52,14 @@ export default defineComponent({
       stagedItems: {} as Record<string, DatasetMeta>,
       trainingOutputName: '',
       selectedTrainingConfig: 'foo.whatever',
+      fineTuneTraining: false,
+      selectedFineTune: null as null | string,
       trainingConfigurations: {
-        configs: [],
-        default: '',
+        training: {
+          configs: [],
+          default: '',
+        },
+        models: {},
       } as TrainingConfigs,
       annotatedFramesOnly: false,
     });
@@ -90,7 +95,16 @@ export default defineComponent({
     onBeforeMount(async () => {
       const configs = await getTrainingConfigurations();
       data.trainingConfigurations = configs;
-      data.selectedTrainingConfig = configs.default;
+      data.selectedTrainingConfig = configs.training.default;
+    });
+
+    const modelNames = computed(() => {
+      if (data.trainingConfigurations.models) {
+        const list = Object.entries(data.trainingConfigurations.models)
+          .map(([, value]) => value.name);
+        return list;
+      }
+      return [];
     });
 
     function toggleStaged(meta: DatasetMeta) {
@@ -172,12 +186,20 @@ export default defineComponent({
     }
 
     async function runTrainingOnFolder() {
+      // Get the full data for fine tuning
+      let foundTrainingModel;
+      if (data.fineTuneTraining) {
+        foundTrainingModel = Object.values(data.trainingConfigurations.models)
+          .find((item) => item.name === data.selectedFineTune);
+      }
       try {
         await runTraining(
           stagedItems.value.map(({ id }) => id),
           data.trainingOutputName,
           data.selectedTrainingConfig,
           data.annotatedFramesOnly,
+          undefined,
+          foundTrainingModel,
         );
         router.push({ name: 'jobs' });
       } catch (err) {
@@ -204,6 +226,7 @@ export default defineComponent({
       nameRules,
       itemsPerPageOptions,
       clientSettings,
+      modelNames,
       models: {
         items: trainedModels,
         headers: trainedHeadersTmpl.concat({
@@ -255,7 +278,10 @@ export default defineComponent({
       <v-card-text>
         Add datasets to the staging area and choose a training configuration.
       </v-card-text>
-      <v-row class="mt-4 pt-0">
+      <v-row
+        class="mt-4 pt-0"
+        dense
+      >
         <v-col sm="5">
           <v-text-field
             v-model="data.trainingOutputName"
@@ -272,7 +298,7 @@ export default defineComponent({
             outlined
             dense
             label="Configuration File (Required)"
-            :items="data.trainingConfigurations.configs"
+            :items="data.trainingConfigurations.training.configs"
             :hint="data.selectedTrainingConfig"
             persistent-hint
           >
@@ -321,6 +347,20 @@ export default defineComponent({
           Train on ({{ staged.items.value.length }}) Datasets
         </v-btn>
       </div>
+      <div class="d-flex flex-row mt-7">
+        <v-checkbox
+          v-model="data.fineTuneTraining"
+          label="Fine Tuning"
+          hint="Fine tune an existing model"
+        />
+        <v-spacer />
+        <v-select
+          v-if="data.fineTuneTraining"
+          v-model="data.selectedFineTune"
+          :items="modelNames"
+          label="Fine Tune Model"
+        />
+      </div>
     </div>
     <div>
       <v-card-title class="text-h4">

client/platform/web-girder/api/rpc.service.ts

@@ -25,8 +25,14 @@ function runTraining(
   config: string,
   annotatedFramesOnly: boolean,
   labelText?: string,
+  fineTuneModel?: {
+    name: string;
+    type: string;
+    path?: string;
+    folderId?: string;
+  },
 ) {
-  return girderRest.post('dive_rpc/train', { folderIds, labelText }, {
+  return girderRest.post('dive_rpc/train', { folderIds, labelText, fineTuneModel }, {
     params: {
       pipelineName, config, annotatedFramesOnly,
     },

client/platform/web-girder/views/RunTrainingMenu.vue

@@ -38,6 +38,8 @@ export default defineComponent({
     const trainingConfigurations = ref<TrainingConfigs | null>(null);
     const selectedTrainingConfig = ref<string | null>(null);
     const annotatedFramesOnly = ref<boolean>(false);
+    const fineTuning = ref<boolean>(false);
+    const selectedFineTune = ref<string>('');
     const {
       request: _runTrainingRequest,
       reset: dismissJobDialog,
@@ -46,10 +48,32 @@ export default defineComponent({
 
     const successMessage = computed(() => `Started training on ${props.selectedDatasetIds.length} dataset(s)`);
 
+    const fineTuneModelList = computed(() => {
+      const modelList: {text: string, type: 'user' | 'system', name: string}[] = [];
+      if (trainingConfigurations.value?.models) {
+        Object.entries(trainingConfigurations.value.models)
+          .forEach(([, value]) => {
+            modelList.push({
+              text: `${value.name} - ${value.folderId ? 'User' : 'System'} Model`,
+              type: value.folderId ? 'user' : 'system',
+              name: value.name,
+            });
+          });
+      }
+      modelList.sort((a, b) => b.type.localeCompare(a.type));
+      return modelList;
+    });
+    const selectedFineTuneObject = computed(() => {
+      if (selectedFineTune.value !== '' && trainingConfigurations.value?.models) {
+        return Object.values(trainingConfigurations.value.models)
+          .find((model) => model.name === selectedFineTune.value);
+      }
+      return undefined;
+    });
     onBeforeMount(async () => {
       const resp = await getTrainingConfigurations();
       trainingConfigurations.value = resp;
-      selectedTrainingConfig.value = resp.default;
+      selectedTrainingConfig.value = resp.training.default;
     });
 
     const trainingDisabled = computed(() => props.selectedDatasetIds.length === 0);
@@ -74,13 +98,16 @@ export default defineComponent({
             selectedTrainingConfig.value,
             annotatedFramesOnly.value,
             labelText.value,
+            selectedFineTuneObject.value,
           );
         }
         return runTraining(
           props.selectedDatasetIds,
           outputPipelineName,
           selectedTrainingConfig.value,
           annotatedFramesOnly.value,
+          undefined,
+          selectedFineTuneObject.value,
         );
       });
       menuOpen.value = false;
@@ -116,6 +143,10 @@ export default defineComponent({
       labelFile,
       clearLabelText,
       simplifyTrainingName,
+      // Fine-Tuning
+      fineTuning,
+      fineTuneModelList,
+      selectedFineTune,
     };
   },
 });
@@ -161,6 +192,7 @@ export default defineComponent({
         <v-card
           v-if="trainingConfigurations"
           outlined
+          class="training-menu"
         >
           <v-card-title class="pb-1">
             Run Training
@@ -198,7 +230,7 @@ export default defineComponent({
               outlined
               class="my-4"
               label="Configuration File"
-              :items="trainingConfigurations.configs"
+              :items="trainingConfigurations.training.configs"
               :hint="selectedTrainingConfig"
               persistent-hint
             >
@@ -225,6 +257,25 @@ export default defineComponent({
               persistent-hint
               class="pt-0"
             />
+            <v-checkbox
+              v-model="fineTuning"
+              label="Fine Tune Model"
+              hint="Fine Tune an existing model"
+              persistent-hint
+              class="pt-0"
+            />
+            <v-select
+              v-if="fineTuning"
+              v-model="selectedFineTune"
+              outlined
+              class="my-4"
+              label="Fine Tune Model"
+              :items="fineTuneModelList"
+              item-value="name"
+              item-text="text"
+              hint="Model to Fine Tune"
+              persistent-hint
+            />
             <v-btn
               depressed
               block
@@ -248,3 +299,10 @@ export default defineComponent({
     />
   </div>
 </template>
+
+<style lang="css" scoped>
+.training-menu {
+  max-height: 90vh;
+  overflow-y: auto;
+}
+</style>

docker-compose.yml

@@ -125,7 +125,7 @@ services:
         reservations:
           devices:
           - driver: nvidia
-            device_ids: ['1']
+            device_ids: ['0']
             capabilities: [gpu]
     environment:
       - "WORKER_WATCHING_QUEUES=training"