update training

Author: choyingw

backbone_nets/mobilenetv2_backbone.py

@@ -1,7 +1,6 @@
 import torch
 from torch import nn
-#from .utils import load_state_dict_from_url
-
+from torch.hub import load_state_dict_from_url
 
 __all__ = ['MobileNetV2', 'mobilenet_v2']
 
@@ -205,5 +204,5 @@ def mobilenet_v2(pretrained=False, progress=True, **kwargs):
     if pretrained:
         state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
                                               progress=progress)
-        model.load_state_dict(state_dict)
+        model.load_state_dict(state_dict, strict=False)
     return model

loss_definition.py

@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-from math import sqrt
 from utils.params import ParamsPack
 param_pack = ParamsPack()
 import math
@@ -33,20 +32,13 @@ def __init__(self):
         super(ParamLoss, self).__init__()
         self.criterion = nn.MSELoss(reduction="none")
 
-    def forward(self, input, target, mode='normal'):
+    def forward(self, input, target, mode = 'normal'):
         if mode == 'normal':
             loss = self.criterion(input[:,:12], target[:,:12]).mean(1) + self.criterion(input[:,12:], target[:,12:]).mean(1)
             return torch.sqrt(loss)
-        if mode == 'no_tex':
-            loss = self.criterion(input[:,:12], target[:,:12]).mean(1) + self.criterion(input[:,12:62], target[:,12:62]).mean(1)
-            return torch.sqrt(loss)
-        if mode == 'only_3dmm':
+        elif mode == 'only_3dmm':
             loss = self.criterion(input[:,:50], target[:,12:62]).mean(1)
             return torch.sqrt(loss)
-        if mode == 'only_pose':
-            loss = self.criterion(input[:,:12], target[:,:12]).mean(1)
-            return loss
-
         return torch.sqrt(loss.mean(1))
 
 

main_train.py

@@ -51,7 +51,6 @@ def parse_args():
     parser.add_argument('--test_initial', default='false', type=str2bool)
     parser.add_argument('--warmup', default=-1, type=int)
     parser.add_argument('--param-fp-train',default='',type=str)
-    parser.add_argument('--data_ver', default='v1', type=str)
     parser.add_argument('--img_size', default=120, type=int)
     parser.add_argument('--save_val_freq', default=10, type=int)
 
@@ -225,7 +224,7 @@ def main():
 
         filename = f'{args.snapshot}_checkpoint_epoch_{epoch}.pth.tar'
         # save checkpoints and current model validation
-        if (epoch % args.save_val_freq == 0) or (epoch==args.epochs) or ((epoch >= 45) and (epoch % 2 ==0)):
+        if (epoch % args.save_val_freq == 0) or (epoch==args.epochs):
             save_checkpoint(
                 {
                     'epoch': epoch,

model_building.py

@@ -33,13 +33,15 @@ def __init__(self, args):
 		self.args = args
 		# backbone definition
 		if 'mobilenet_v2' in self.args.arch:
-			self.backbone = getattr(mobilenetv2_backbone, args.arch)()
+			self.backbone = getattr(mobilenetv2_backbone, args.arch)(pretrained=False)
 		elif 'mobilenet' in self.args.arch:
 			self.backbone = getattr(mobilenetv1_backbone, args.arch)()		
 		elif 'resnet' in self.args.arch:
 			self.backbone = getattr(resnet_backbone, args.arch)(pretrained=False)
 		elif 'ghostnet' in self.args.arch:
 			self.backbone = getattr(ghostnet_backbone, args.arch)()
+		else:
+			raise RuntimeError("Please choose [mobilenet_v2, mobilenet_1, resnet50, or ghostnet]")
 
 	def forward(self,input, target):
 		"""Training time forward"""
@@ -102,7 +104,7 @@ def reconstruct_vertex_62(self, param, whitening=True, dense=False, transform=Tr
 		transform: whether transform to image space
 		Working with batched tensors. Using Fortan-type reshape.
 		"""
-		
+
 		if whitening:
 			if param.shape[1] == 62:
 				param_ = param * self.param_std[:62] + self.param_mean[:62]
@@ -134,16 +136,16 @@ def forward(self, input, target):
 
 		vertex_lmk = self.reconstruct_vertex_62(_3D_attr, dense=False)
 		vertex_GT_lmk = self.reconstruct_vertex_62(_3D_attr_GT, dense=False)
-		self.loss['loss_LMK_f0'] = 0.01 *self.LMKLoss_3D(vertex_lmk, vertex_GT_lmk, kp=True)		
+		self.loss['loss_LMK_f0'] = 0.05 *self.LMKLoss_3D(vertex_lmk, vertex_GT_lmk, kp=True)		
 		self.loss['loss_Param_In'] = 0.02 * self.ParamLoss(_3D_attr, _3D_attr_GT)
-		
+
 		point_residual = self.forwardDirection(vertex_lmk, avgpool, _3D_attr[:,12:52], _3D_attr[:,52:62])
-		vertex_lmk = vertex_lmk + 0.01 * point_residual
-		self.loss['loss_LMK_pointNet'] = 0.01 * self.LMKLoss_3D(vertex_lmk, vertex_GT_lmk, kp=True)
+		vertex_lmk = vertex_lmk + 0.05 * point_residual
+		self.loss['loss_LMK_pointNet'] = 0.05 * self.LMKLoss_3D(vertex_lmk, vertex_GT_lmk, kp=True)
 
 		_3D_attr_S2 = self.reverseDirection(vertex_lmk)
 		self.loss['loss_Param_S2'] = 0.02 * self.ParamLoss(_3D_attr_S2, _3D_attr_GT, mode='only_3dmm')
-		self.loss['loss_Param_S1S2'] = 1 * self.ParamLoss(_3D_attr_S2, _3D_attr, mode='only_3dmm')
+		self.loss['loss_Param_S1S2'] = 0.001 * self.ParamLoss(_3D_attr_S2, _3D_attr, mode='only_3dmm')
 
 		return self.loss
 

train_script.sh

@@ -11,10 +11,10 @@ python3 main_train.py --arch="mobilenet_v2" \
     --snapshot="ckpts/SynergyNet" \
     --param-fp-train='./3dmm_data/param_all_norm_v201.pkl' \
     --warmup=5 \
-    --batch-size=350 \
-    --base-lr=0.01 \
-    --epochs=80 \
-    --milestones=48,64 \
+    --batch-size=900 \
+    --base-lr=0.027 \
+    --epochs=50 \
+    --milestones=30,40 \
     --print-freq=50 \
     --devices-id=0 \
     --workers=8 \

utils/ddfa.py

@@ -228,6 +228,7 @@ def __call__(self, img, gt=None):
         if not(_is_tensor_image(img)):
             raise TypeError('img should be tensor. Got {}'.format(type(img)))
         if img.ndim == 3:
+            crop_backgnd[:, crop_margins:h-1*crop_margins, crop_margins:w-1*crop_margins] = img[:, crop_margins: h-crop_margins, crop_margins: w-crop_margins]
             # random center crop
             if (rand < self.prob) and (self.mode=='train'):
                 func = self.switcher.get(random.randint(1,7))