Aktualisieren models.py

models.py

@@ -168,49 +168,56 @@ class ActorOnlySummarisationModel(SummarisationModel):
 
 class SummarisationModelWithCrossEntropyLoss(SummarisationModel):
     
-    def _train(self, dataset, epochs=20, batch_size=20, learning_rate=0.001, shuffle=True):
-        
-        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
-        loss_fn = nn.BCELoss(reduction='sum')
-        
-        for _ in range(epochs):
-            
-            training_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
-            
-            for batch in training_dataloader:
-
-                optimizer.zero_grad()
+    #def _train(self, dataset, epochs=20, batch_size=20, learning_rate=0.001, shuffle=True):
+    #    
+    #    optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
+    #    loss_fn = nn.BCELoss(reduction='sum')
+    #    
+    #    for _ in range(epochs):
+    #        
+    #        training_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
+    #        
+    #        for batch in training_dataloader:
+    #
+    #            optimizer.zero_grad()
+    #
+    #            for datapoint in batch:
+    #                
+    #                _, probs = self.__call__(datapoint.document)
+    #
+    #                loss = loss_fn(probs, datapoint.bin_summary)
+    #
+    #                loss.backward()
+    #            
+    #            optimizer.step()
 
-                for datapoint in batch:
-                    
-                    _, probs = self.__call__(datapoint.document)
-
-                    loss = loss_fn(probs, datapoint.bin_summary)
+        # eval
+    
+    def __init__(self, learning_rate=0.001):
 
-                    loss.backward()
-                
-                optimizer.step()
+        super().__init__()
 
-        # eval
+        self.loss_fn = nn.BCELoss(reduction='sum')
+        self.optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
 
-    def epoch(self, dataloader, optimizer):
+    def epoch(self, dataloader):
+        
         
-        loss_fn = nn.BCELoss(reduction='sum')
         self.train()
             
         for batch in dataloader:
 
-            optimizer.zero_grad()
+            self.optimizer.zero_grad()
 
             for datapoint in batch:
                 
                 _, probs = self.__call__(datapoint.document)
 
-                loss = loss_fn(probs, datapoint.bin_summary)
+                loss = self.loss_fn(probs, datapoint.bin_summary)
 
                 loss.backward()
             
-            optimizer.step()
+            self.optimizer.step()