use pretrain backbone

Author: CoinCheung

configs/bisenetv2.py

@@ -3,7 +3,7 @@
 cfg = dict(
     model_type='bisenetv2',
     num_aux_heads=4,
-    lr_start = 1 * 5e-3,
+    lr_start = 5e-3,
     weight_decay=5e-4,
     warmup_iters = 1000,
     max_iter = 150000,

dist_train.sh

@@ -1,6 +1,6 @@
 
 export CUDA_VISIBLE_DEVICES=6,7
-PORT=52330
+PORT=52332
 NGPUS=2
 
 python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/train_amp.py --model bisenetv2 --port $PORT

lib/models/bisenetv2.py

@@ -2,6 +2,9 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.utils.model_zoo as modelzoo
+
+backbone_url = 'https://github.com/CoinCheung/BiSeNet/releases/download/0.0.0/backbone_v2.pth'
 
 
 class ConvBNReLU(nn.Module):
@@ -325,7 +328,6 @@ def __init__(self, n_classes, output_aux=True):
             self.aux5_4 = SegmentHead(128, 128, n_classes, up_factor=32)
 
         self.init_weights()
-        self.load_pretrain()
 
     def forward(self, x):
         size = x.size()[2:]
@@ -354,13 +356,13 @@ def init_weights(self):
                 else:
                     nn.init.ones_(module.weight)
                 nn.init.zeros_(module.bias)
+        self.load_pretrain()
 
     def load_pretrain(self):
-        state = torch.load('pretrained/bisenetv2_pretrain.pth', map_location='cpu')
-        state = {k:v for k,v in state.items() if not k in ('fc', 'head', 'dense_head')}
+        state = modelzoo.load_url(backbone_url)
         for name, child in self.named_children():
             if name in state.keys():
-                child.load_state_dict(state[name])
+                child.load_state_dict(state[name], strict=True)
 
 
     def get_params(self):