Update modeling_diffusiondet.py

modeling_diffusiondet.py

@@ -127,7 +127,13 @@ class DiffusionDet(PreTrainedModel):
         self.criterion = CriterionDynamicK(config, num_classes=self.num_classes, weight_dict=weight_dict)
 
     def _init_weights(self, module):
-        torch.nn.init.xavier_uniform_(module.weight)
+        if isinstance(module, nn.Linear) or isinstance(module, nn.Conv2d):
+            init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
+        if module.bias is not None:
+            init.constant_(module.bias, 0)
+        elif isinstance(module, nn.BatchNorm2d):
+            init.constant_(module.weight, 1)
+            init.constant_(module.bias, 0)
 
     def predict_noise_from_start(self, x_t, t, x0):
         return (