Implement weight initialization

1 files changed, 14 insertions, 0 deletions

diff --git a/models/model.py b/models/model.py
index c407d6c..ebd6aba 100644
--- a/models/model.py
+++ b/models/model.py
@@ -2,6 +2,7 @@ from typing import Union, Optional
 
 import numpy as np
 import torch
+import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
@@ -67,6 +68,7 @@ class Model:
         self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, 500, 0.9)
 
         self.rbg_pn.train()
+        self.rbg_pn.apply(self.init_weights)
         for iter_i, (x_c1, x_c2) in enumerate(dataloader):
             loss = self.rbg_pn(x_c1['clip'], x_c2['clip'], x_c1['label'])
             loss.backward()
@@ -76,6 +78,18 @@ class Model:
             if iter_i == self.meta['total_iter']:
                 break
 
+    @staticmethod
+    def init_weights(m):
+        if isinstance(m, nn.modules.conv._ConvNd):
+            nn.init.xavier_uniform_(m.weight)
+        elif isinstance(m, nn.modules.batchnorm._NormBase):
+            nn.init.normal_(m.weight, 1.0, 0.01)
+            nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+        elif isinstance(m, nn.Parameter):
+            nn.init.xavier_uniform_(m)
+
     def _parse_dataset_config(
             self,
             dataset_config: DatasetConfiguration