update seed

Author: bubbliiiing

nets/attention.py

@@ -97,16 +97,24 @@ def __init__(self, channel, reduction=16):
         self.sigmoid_w = nn.Sigmoid()
 
     def forward(self, x):
+        # batch_size, c, h, w
         _, _, h, w = x.size()
 
+        # batch_size, c, h, w => batch_size, c, h, 1 => batch_size, c, 1, h
         x_h = torch.mean(x, dim = 3, keepdim = True).permute(0, 1, 3, 2)
+        # batch_size, c, h, w => batch_size, c, 1, w
         x_w = torch.mean(x, dim = 2, keepdim = True)
- 
+
+        # batch_size, c, 1, w cat batch_size, c, 1, h => batch_size, c, 1, w + h
+        # batch_size, c, 1, w + h => batch_size, c / r, 1, w + h
         x_cat_conv_relu = self.relu(self.bn(self.conv_1x1(torch.cat((x_h, x_w), 3))))
- 
+
+        # batch_size, c / r, 1, w + h => batch_size, c / r, 1, h and batch_size, c / r, 1, w
         x_cat_conv_split_h, x_cat_conv_split_w = x_cat_conv_relu.split([h, w], 3)
- 
+
+        # batch_size, c / r, 1, h => batch_size, c / r, h, 1 => batch_size, c, h, 1
         s_h = self.sigmoid_h(self.F_h(x_cat_conv_split_h.permute(0, 1, 3, 2)))
+        # batch_size, c / r, 1, w => batch_size, c, 1, w
         s_w = self.sigmoid_w(self.F_w(x_cat_conv_split_w))
 
         out = x * s_h.expand_as(x) * s_w.expand_as(x)

requirements.txt

@@ -1,9 +1,10 @@
+torch
+torchvision
+tensorboard
 scipy==1.2.1
 numpy==1.17.0
 matplotlib==3.1.2
 opencv_python==4.1.2.30
-torch==1.2.0
-torchvision==0.4.0
 tqdm==4.60.0
 Pillow==8.2.0
-h5py==2.10.0
+h5py==2.10.0

train.py

@@ -3,6 +3,7 @@
 #-------------------------------------#
 import datetime
 import os
+from functools import partial
 
 import numpy as np
 import torch
@@ -18,7 +19,8 @@
                                 weights_init)
 from utils.callbacks import EvalCallback, LossHistory
 from utils.dataloader import YoloDataset, yolo_dataset_collate
-from utils.utils import get_anchors, get_classes, show_config
+from utils.utils import (get_anchors, get_classes, seed_everything,
+                         show_config, worker_init_fn)
 from utils.utils_fit import fit_one_epoch
 
 '''
@@ -43,6 +45,11 @@
     #           没有GPU可以设置成False
     #---------------------------------#
     Cuda = True
+    #----------------------------------------------#
+    #   Seed    用于固定随机种子
+    #           使得每次独立训练都可以获得一样的结果
+    #----------------------------------------------#
+    seed = 11
     #---------------------------------------------------------------------#
     #   distributed     用于指定是否使用单机多卡分布式运行
     #                   终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
@@ -265,6 +272,7 @@
     train_annotation_path   = '2007_train.txt'
     val_annotation_path     = '2007_val.txt'
 
+    seed_everything(seed)
     #------------------------------------------------------#
     #   设置用到的显卡
     #------------------------------------------------------#
@@ -280,6 +288,7 @@
     else:
         device          = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         local_rank      = 0
+        rank            = 0
 
     #----------------------------------------------------#
     #   获取classes和anchor
@@ -482,9 +491,11 @@
             shuffle         = True
 
         gen             = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
-                                    drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler)
+                                    drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler, 
+                                    worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
         gen_val         = DataLoader(val_dataset  , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True, 
-                                    drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler)
+                                    drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler, 
+                                    worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
 
         #----------------------#
         #   记录eval的map曲线
@@ -533,9 +544,11 @@
                     batch_size = batch_size // ngpus_per_node
 
                 gen             = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
-                                            drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler)
+                                            drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler, 
+                                            worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
                 gen_val         = DataLoader(val_dataset  , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True, 
-                                            drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler)
+                                            drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler, 
+                                            worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
 
                 UnFreeze_flag = True
 

utils/utils.py

@@ -1,4 +1,7 @@
+import random
+
 import numpy as np
+import torch
 from PIL import Image
 
 #---------------------------------------------------------#
@@ -57,6 +60,27 @@ def get_lr(optimizer):
     for param_group in optimizer.param_groups:
         return param_group['lr']
 
+#---------------------------------------------------#
+#   设置种子
+#---------------------------------------------------#
+def seed_everything(seed=11):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+#---------------------------------------------------#
+#   设置Dataloader的种子
+#---------------------------------------------------#
+def worker_init_fn(worker_id, rank, seed):
+    worker_seed = rank + seed
+    random.seed(worker_seed)
+    np.random.seed(worker_seed)
+    torch.manual_seed(worker_seed)
+
 def preprocess_input(image):
     image /= 255.0
     return image