distribute bert

Author: Wonters

src/mixins.py

@@ -70,7 +70,8 @@ def _train_batch(self, x, y):
         inputs = self.tokenizer(x, return_tensors="pt", truncation=True, padding=True)
         if isinstance(inputs, dict) and inputs["input_ids"]:
             inputs["input_ids"] = inputs["input_ids"].float()
-        if False and isinstance(y, torch.Tensor) and y.dtype == torch.float32:
+        # todo : fix this case for bert vs lstm
+        if True and isinstance(y, torch.Tensor) and y.dtype == torch.float32:
             labels = y.long()
         else:
             labels = y.float()
@@ -84,6 +85,7 @@ def _train_batch(self, x, y):
             _, preds = torch.max(outputs.logits, dim=1)
         except AttributeError:
             loss = self.criterion(outputs, labels)
+            # todo : error on accuracy return always 0 
             preds = outputs 
         correct = (preds == labels).sum().item()    
         acc = correct / len(labels)

src/ml.py

@@ -426,7 +426,11 @@ def load_checkpoint(self):
         self.criterion = torch.nn.CrossEntropyLoss()
 
     def save(self):
-        self.model.save_pretrained(self.checkpoint)
+        print(dist.is_available(), dist.is_initialized())
+        if dist.is_available() and dist.is_initialized():
+            self.model.module.save_pretrained(self.checkpoint)
+        else:
+            self.model.save_pretrained(self.checkpoint)
         self.tokenizer.save_pretrained(self.checkpoint)
 
     def predict(self, x: list):