cc

configs/experiment/miniagent-bert-mlp-abs_diff-mult.yaml

@@ -0,0 +1,37 @@
+# @package _global_
+
+defaults:
+  - override /data: mixed
+  - override /model: miniagent
+  - override /callbacks: default
+  - override /trainer: gpu
+
+seed: 42
+
+model:
+  lr: 0.001
+  bert_model: bert-base-uncased
+
+  inst_proj_model:
+    _target_: src.models.mlp_module.MLPProjection
+    input_dim: 768
+    hidden_dim: 768
+    output_dim: 768
+
+  tool_proj_model:
+    _target_: src.models.mlp_module.MLPProjection
+    input_dim: 768
+    hidden_dim: 768
+    output_dim: 768
+
+  pred_model:
+    _target_: src.models.mlp_module.MLPPrediction
+    input_dim: 768
+    use_abs_diff: true
+    use_mult: true
+
+data:
+  bert_model: bert-base-uncased
+  seed: 42
+  batch_size: 128
+  tool_capacity: 16

configs/trainer/default.yaml

@@ -8,6 +8,8 @@ max_epochs: 10
 accelerator: cpu
 devices: 1
 
+log_every_n_steps: 10
+
 # mixed precision for extra speed-up
 # precision: 16
 

src/data/mixed_datamodule.py

@@ -47,6 +47,7 @@ class MixedDataModule(LightningDataModule):
             batch_size=self.batch_size,
             shuffle=True,
             num_workers=self.num_workers,
+            drop_last=True,
         )
 
     def val_dataloader(self):

src/models/miniagent_module.py

@@ -65,10 +65,12 @@ class MiniAgentModule(LightningModule):
         pred = self.pred_model(inst_emb_r, tool_emb_r)  # [BxB, 1]
         pred = pred.view(B, B)  # [B, B]
 
-        target = torch.eye(B, device=pred.device).float()
+        # target = torch.eye(B, device=pred.device).float()
 
-        pos_weight = torch.tensor([B - 1], device=pred.device)
-        loss = F.binary_cross_entropy_with_logits(pred, target, pos_weight=pos_weight)
+        # pos_weight = torch.tensor([B - 1], device=pred.device)
+        # loss = F.binary_cross_entropy_with_logits(pred, target, pos_weight=pos_weight)
+        labels = torch.arange(B, device=pred.device).long()
+        loss = (F.cross_entropy(pred, labels) + F.cross_entropy(pred.T, labels)) * 0.5
 
         self.log("train/loss", loss, on_step=True, sync_dist=True, prog_bar=True)