From 9095549b46bf0f371340b040bda172fbc2998014 Mon Sep 17 00:00:00 2001
From: Fabio <ffd2011@nyu.edu>
Date: Wed, 20 Nov 2024 14:23:29 -0500
Subject: [PATCH] Adjust hidden state of training

---
 config/A8_swin.yaml     | 11 +++++++++++
 config/M2.yaml          |  5 +++++
 config/M3.yaml          |  7 ++++++-
 config/M4_swin.yaml     | 11 +++++++++++
 config/M5.yaml          |  5 +++++
 config/R16_swin.yaml    | 11 +++++++++++
 config/R18.yaml         |  5 +++++
 config/R19_swin.yaml    | 10 ++++++++++
 scripts/omnimix.sh      |  6 +++---
 step_recog/iterators.py | 11 ++---------
 step_recog/models.py    |  2 +-
 11 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/config/A8_swin.yaml b/config/A8_swin.yaml
index 4ef02bb..581b025 100644
--- a/config/A8_swin.yaml
+++ b/config/A8_swin.yaml
@@ -2,3 +2,14 @@ _BASE_: A8.yaml
 MODEL:
   CLASS: OmniTransformer_v4
   OMNIGRU_CHECKPOINT_URL: /home/user/models/A8_swin.pt
+
+DATASET:
+  CLASS: 'Milly_multifeature_v6'  
+
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.0001
+  LINEAR_WARMUP_EPOCHS: 3 #default None
+  WEIGHT_DECAY: 0.05
+  SCHEDULER: 'cos'
+  EPOCHS: 30
\ No newline at end of file
diff --git a/config/M2.yaml b/config/M2.yaml
index cb67856..021a6b9 100644
--- a/config/M2.yaml
+++ b/config/M2.yaml
@@ -12,6 +12,11 @@ DATASET:
 OUTPUT:
   LOCATION: "/home/user/output"
 
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.001
+  WEIGHT_DECAY: 0.0001
+
 SKILLS:
   - NAME: M2 - Apply Tourniquet
     STEPS:
diff --git a/config/M3.yaml b/config/M3.yaml
index 3662375..6a62f74 100644
--- a/config/M3.yaml
+++ b/config/M3.yaml
@@ -11,7 +11,12 @@ DATASET:
 
 OUTPUT:
   LOCATION: "/home/user/output"  
-  
+
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.001
+  WEIGHT_DECAY: 0.001
+
 SKILLS:
   - NAME: M3 - Apply pressure dressing
     STEPS:
diff --git a/config/M4_swin.yaml b/config/M4_swin.yaml
index fc50524..be78ab3 100644
--- a/config/M4_swin.yaml
+++ b/config/M4_swin.yaml
@@ -2,3 +2,14 @@ _BASE_: M4.yaml
 MODEL:
   CLASS: OmniTransformer_v3
   OMNIGRU_CHECKPOINT_URL: /home/user/models/M4_swin.pt
+
+DATASET:
+  CLASS: 'Milly_multifeature_v6'  
+
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.0001
+  LINEAR_WARMUP_EPOCHS: 3 #default None
+  WEIGHT_DECAY: 0.05
+  SCHEDULER: 'cos'
+  EPOCHS: 30
\ No newline at end of file
diff --git a/config/M5.yaml b/config/M5.yaml
index db2d7fe..eab3e07 100644
--- a/config/M5.yaml
+++ b/config/M5.yaml
@@ -12,6 +12,11 @@ DATASET:
 OUTPUT:
   LOCATION: "/home/user/output"  
 
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.001
+  WEIGHT_DECAY: 0.00001
+
 SKILLS:
   - NAME: M5 - X-Stat
     STEPS:
diff --git a/config/R16_swin.yaml b/config/R16_swin.yaml
index f0c681a..87272dc 100644
--- a/config/R16_swin.yaml
+++ b/config/R16_swin.yaml
@@ -2,3 +2,14 @@ _BASE_: R16.yaml
 MODEL:
   CLASS: OmniTransformer_v3
   OMNIGRU_CHECKPOINT_URL: /home/user/models/R16_swin.pt
+
+DATASET:
+  CLASS: 'Milly_multifeature_v6'  
+
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.0001
+  LINEAR_WARMUP_EPOCHS: 3 #default None
+  WEIGHT_DECAY: 0.05
+  SCHEDULER: 'cos'
+  EPOCHS: 30
\ No newline at end of file
diff --git a/config/R18.yaml b/config/R18.yaml
index e2357c3..8b30030 100644
--- a/config/R18.yaml
+++ b/config/R18.yaml
@@ -12,6 +12,11 @@ DATASET:
 OUTPUT:
   LOCATION: "/home/user/output"    
 
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.001
+  WEIGHT_DECAY: 0.01
+
 SKILLS:
   - NAME: R18 - Apply chest seal
     STEPS:
diff --git a/config/R19_swin.yaml b/config/R19_swin.yaml
index 08ae5e2..7151017 100644
--- a/config/R19_swin.yaml
+++ b/config/R19_swin.yaml
@@ -3,3 +3,13 @@ MODEL:
   CLASS: OmniTransformer_v3
   OMNIGRU_CHECKPOINT_URL: /home/user/models/R19_swin.pt
 
+DATASET:
+  CLASS: 'Milly_multifeature_v6'  
+
+TRAIN:
+  OPT: "adamw"  #adam adamw sgd rmsprop
+  LR: 0.0001
+  LINEAR_WARMUP_EPOCHS: 3 #default None
+  WEIGHT_DECAY: 0.05
+  SCHEDULER: 'cos'
+  EPOCHS: 30
\ No newline at end of file
diff --git a/scripts/omnimix.sh b/scripts/omnimix.sh
index fcae475..6da5ee4 100644
--- a/scripts/omnimix.sh
+++ b/scripts/omnimix.sh
@@ -41,7 +41,7 @@ sbatch <<EOSBATCH
 #SBATCH --mem 64GB
 #SBATCH --time 2-00:00:00
 #SBATCH --gres gpu:1
-#SBATCH --job-name step-recog-k$KFOLD_ITER$DESC
+#SBATCH --job-name step-k$KFOLD_ITER$DESC
 #SBATCH --output logs/%J_step-k$KFOLD_ITER$DESC.out
 #SBATCH --mail-type=BEGIN,END,ERROR
 #SBATCH --mail-user=$USER@nyu.edu
@@ -70,8 +70,8 @@ sbatch <<EOSBATCH
 #SBATCH --mem 64GB
 #SBATCH --time 2-00:00:00
 #SBATCH --gres gpu:1
-#SBATCH --job-name step$DESC
-#SBATCH --output logs/%J_step-recog$DESC.out
+#SBATCH --job-name step-$DESC
+#SBATCH --output logs/%J_step-$DESC.out
 #SBATCH --mail-type=BEGIN,END,ERROR
 #SBATCH --mail-user=$USER@nyu.edu
 #SBATCH --account=$PROJECT_ACCOUNT
diff --git a/step_recog/iterators.py b/step_recog/iterators.py
index 7693ee0..c6e0cfe 100644
--- a/step_recog/iterators.py
+++ b/step_recog/iterators.py
@@ -71,7 +71,6 @@ def build_optimizer(model, cfg):
 def train_step_GRU(model, criterion, criterion_t, optimizer, loader, is_training, device, cfg, progress, class_weight):
   if is_training:
     model.train()
-    h = model.init_hidden(cfg.TRAIN.BATCH_SIZE)
   else:
     model.eval()  
 
@@ -94,13 +93,9 @@ def train_step_GRU(model, criterion, criterion_t, optimizer, loader, is_training
     action, obj, frame, audio, label, label_t, mask, _, _ = next(loader_iterator)
     label = nn.functional.one_hot(label, model.number_classes)
 
-    if not is_training:
-      h = model.init_hidden(len(label))
-
-    h = torch.zeros_like(h)
     optimizer.zero_grad()
 
-    out, h = model(action.to(device).float(), h, audio.to(device).float(), obj.to(device).float(), frame.to(device).float(), return_last_step = False)
+    out, _ = model(action=action.to(device).float(), aud=audio.to(device).float(), objs=obj.to(device).float(), frame=frame.to(device).float(), return_last_step = False)
 
     mask    = mask.to(out.device)  
     label   = label.to(out.device)
@@ -468,9 +463,7 @@ def evaluate_GRU(model, data_loader, cfg):
 #  ipdb.set_trace()
 
   for action, obj, frame, audio, label, _, mask, frame_idx, videos in data_loader:
-    h = model.init_hidden(len(action))
-
-    out, _ = model(action.to(device).float(), h, audio.to(device).float(), obj.to(device).float(), frame.to(device).float(), return_last_step = False)
+    out, _ = model(action=action.to(device).float(), aud=audio.to(device).float(), objs=obj.to(device).float(), frame=frame.to(device).float(), return_last_step = False)
     out    = torch.softmax(out[..., :model.number_classes], dim = -1).cpu().detach().numpy()
     label  = label.cpu().numpy()
     frame_idx = frame_idx.cpu().numpy()
diff --git a/step_recog/models.py b/step_recog/models.py
index b2e2b6a..cdfd47e 100644
--- a/step_recog/models.py
+++ b/step_recog/models.py
@@ -112,7 +112,7 @@ def forward(self, action, h=None, aud=None, objs=None, frame=None, return_last_s
             x.append(obj_in)
 
         x = torch.concat(x, -1) if len(x) > 1 else x[0]            
-        out, h = self.gru(x, h)
+        out, h = self.gru(x, self.init_hidden(x.shape[0]) if h is None else h)
         out = self.relu(out[:, -1]) if return_last_step else self.relu(out)
         out = self.fc(out)
         return out, h