FER alpha 0.1

.gitignore

@@ -1,4 +1,5 @@
 FER/models/__pycache__
 FER/__pycache__
 .env
-.venv
+.venv
+FER/Images/

FER/data_preprocessing/sam.py

@@ -0,0 +1,63 @@
+import torch
+
+
+class SAM(torch.optim.Optimizer):
+    def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs):
+        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
+
+        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
+        super(SAM, self).__init__(params, defaults)
+
+        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        self.param_groups = self.base_optimizer.param_groups
+
+    @torch.no_grad()
+    def first_step(self, zero_grad=False):
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = group["rho"] / (grad_norm + 1e-12)
+
+            for p in group["params"]:
+                if p.grad is None: continue
+                self.state[p]["old_p"] = p.data.clone()
+                e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale.to(p)
+                p.add_(e_w)  # climb to the local maximum "w + e(w)"
+
+        if zero_grad: self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad=False):
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None: continue
+                p.data = self.state[p]["old_p"]  # get back to "w" from "w + e(w)"
+
+        self.base_optimizer.step()  # do the actual "sharpness-aware" update
+
+        if zero_grad: self.zero_grad()
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
+        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass
+
+        self.first_step(zero_grad=True)
+        closure()
+        self.second_step()
+
+    def _grad_norm(self):
+        shared_device = self.param_groups[0]["params"][0].device  # put everything on the same device, in case of model parallelism
+        norm = torch.norm(
+                    torch.stack([
+                        ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
+                        for group in self.param_groups for p in group["params"]
+                        if p.grad is not None
+                    ]),
+                    p=2
+               )
+        return norm
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        self.base_optimizer.param_groups = self.param_groups
+

FER/detectfaces.py

@@ -4,13 +4,10 @@ import torch
 import os
 import time
 from PIL import Image
+from main import RecorderMeter1, RecorderMeter
 
 # Define the path to the model checkpoint
-# Get the directory of the current file (models/PosterV2_7cls.py)
-script_dir = os.path.dirname(os.path.abspath(__file__))
-
-# Construct the full path to the model file
-model_path = os.path.join(script_dir, r"models\checkpoints\raf-db-model_best.pth")
+model_path = os.path.abspath(r"FER\models\checkpoints\raf-db-model_best.pth")
 
 # Determine the available device for model execution
 if torch.backends.mps.is_available():
@@ -39,7 +36,7 @@ def main():
     if model_path is not None:
         if os.path.isfile(model_path):
             print("=> loading checkpoint '{}'".format(model_path))
-            checkpoint = torch.load(model_path, map_location=device)
+            checkpoint = torch.load(model_path, map_location=device, weights_only=False)
             best_acc = checkpoint["best_acc"]
             best_acc = best_acc.to()
             print(f"best_acc:{best_acc}")
@@ -50,7 +47,9 @@ def main():
                 )
             )
         else:
-            print("=> no checkpoint found at '{}'".format(model_path))
+            print(
+                "[!] detectfaces.py => no checkpoint found at '{}'".format(model_path)
+            )
         # Start webcam capture and prediction
         imagecapture(model)
         return
@@ -92,7 +91,7 @@ def imagecapture(model):
 
         # If faces are detected, proceed with prediction
         if len(faces) > 0:
-            currtimeimg = time.strftime("%H:%M:%S:%.*f")
+            currtimeimg = time.strftime("%H:%M:%S")
             print(f"[!]Face detected at {currtimeimg}")
             # Crop the face region
             face_region = frame[
@@ -105,12 +104,12 @@ def imagecapture(model):
             )
             print("[!]Start Expressions")
             # Record the prediction start time
-            starttime = time.strftime("%H:%M:%S:%.*f")
+            starttime = time.strftime("%H:%M:%S")
             print(f"-->Prediction starting at {starttime}")
             # Perform emotion prediction
             predict(model, image_path=face_pil_image)
             # Record the prediction end time
-            endtime = time.strftime("%H:%M:%S:%.*f")
+            endtime = time.strftime("%H:%M:%S")
             print(f"-->Done prediction at {endtime}")
 
             # Stop capturing once prediction is complete

FER/main.py

@@ -21,7 +21,7 @@ import torchvision.transforms as transforms
 import numpy as np
 import datetime
 from torchsampler import ImbalancedDatasetSampler
-from models.PosterV2_7cls import *
+from models.PosterV2_7cls import pyramid_trans_expr2
 
 
 warnings.filterwarnings("ignore", category=UserWarning)

FER/models/PosterV2_7cls.py

@@ -5,7 +5,7 @@ from torch.nn import functional as F
 from .mobilefacenet import MobileFaceNet
 from .ir50 import Backbone
 from .vit_model import VisionTransformer, PatchEmbed
-from timm.models.layers import trunc_normal_, DropPath
+from timm.layers import trunc_normal_, DropPath
 from thop import profile
 
 
@@ -315,6 +315,7 @@ class pyramid_trans_expr2(nn.Module):
         face_landback_checkpoint = torch.load(
             mobilefacenet_path,
             map_location=lambda storage, loc: storage,
+            weights_only=False,
         )
         self.face_landback.load_state_dict(face_landback_checkpoint["state_dict"])
 
@@ -325,8 +326,7 @@ class pyramid_trans_expr2(nn.Module):
 
         self.ir_back = Backbone(50, 0.0, "ir")
         ir_checkpoint = torch.load(
-            ir50_path,
-            map_location=lambda storage, loc: storage,
+            ir50_path, map_location=lambda storage, loc: storage, weights_only=False
         )
 
         self.ir_back = load_pretrained_weights(self.ir_back, ir_checkpoint)

FER/models/PosterV2_8cls.py

@@ -4,13 +4,15 @@ from torch.nn import functional as F
 from .mobilefacenet import MobileFaceNet
 from .ir50 import Backbone
 from .vit_model_8 import VisionTransformer, PatchEmbed
-from timm.models.layers import trunc_normal_, DropPath
+from timm.layers import trunc_normal_, DropPath
 from thop import profile
 
+
 def load_pretrained_weights(model, checkpoint):
     import collections
-    if 'state_dict' in checkpoint:
-        state_dict = checkpoint['state_dict']
+
+    if "state_dict" in checkpoint:
+        state_dict = checkpoint["state_dict"]
     else:
         state_dict = checkpoint
     model_dict = model.state_dict()
@@ -19,7 +21,7 @@ def load_pretrained_weights(model, checkpoint):
     for k, v in state_dict.items():
         # If the pretrained state_dict was saved as nn.DataParallel,
         # keys would contain "module.", which should be ignored.
-        if k.startswith('module.'):
+        if k.startswith("module."):
             k = k[7:]
         if k in model_dict and model_dict[k].size() == v.size():
             new_state_dict[k] = v
@@ -30,9 +32,10 @@ def load_pretrained_weights(model, checkpoint):
     model_dict.update(new_state_dict)
 
     model.load_state_dict(model_dict)
-    print('load_weight', len(matched_layers))
+    print("load_weight", len(matched_layers))
     return model
 
+
 def window_partition(x, window_size, h_w, w_w):
     """
     Args:
@@ -44,14 +47,18 @@ def window_partition(x, window_size, h_w, w_w):
     """
     B, H, W, C = x.shape
     x = x.view(B, h_w, window_size, w_w, window_size, C)
-    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    windows = (
+        x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    )
     return windows
 
+
 class window(nn.Module):
     def __init__(self, window_size, dim):
         super(window, self).__init__()
         self.window_size = window_size
         self.norm = nn.LayerNorm(dim)
+
     def forward(self, x):
         x = x.permute(0, 2, 3, 1)
         B, H, W, C = x.shape
@@ -63,21 +70,23 @@ class window(nn.Module):
         x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
         return x_windows, shortcut
 
+
 class WindowAttentionGlobal(nn.Module):
     """
     Global window attention based on: "Hatamizadeh et al.,
     Global Context Vision Transformers <https://arxiv.org/abs/2206.09959>"
     """
 
-    def __init__(self,
-                 dim,
-                 num_heads,
-                 window_size,
-                 qkv_bias=True,
-                 qk_scale=None,
-                 attn_drop=0.,
-                 proj_drop=0.,
-                 ):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
         """
         Args:
             dim: feature size dimension.
@@ -94,9 +103,10 @@ class WindowAttentionGlobal(nn.Module):
         self.window_size = window_size
         self.num_heads = num_heads
         head_dim = torch.div(dim, num_heads)
-        self.scale = qk_scale or head_dim ** -0.5
+        self.scale = qk_scale or head_dim**-0.5
         self.relative_position_bias_table = nn.Parameter(
-            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
+        )
         coords_h = torch.arange(self.window_size[0])
         coords_w = torch.arange(self.window_size[1])
         coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
@@ -112,7 +122,7 @@ class WindowAttentionGlobal(nn.Module):
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
-        trunc_normal_(self.relative_position_bias_table, std=.02)
+        trunc_normal_(self.relative_position_bias_table, std=0.02)
         self.softmax = nn.Softmax(dim=-1)
 
     def forward(self, x, q_global):
@@ -122,14 +132,23 @@ class WindowAttentionGlobal(nn.Module):
         B = q_global.shape[0]
         head_dim = int(torch.div(C, self.num_heads).item())
         B_dim = int(torch.div(B_, B).item())
-        kv = self.qkv(x).reshape(B_, N, 2, self.num_heads, head_dim).permute(2, 0, 3, 1, 4)
+        kv = (
+            self.qkv(x)
+            .reshape(B_, N, 2, self.num_heads, head_dim)
+            .permute(2, 0, 3, 1, 4)
+        )
         k, v = kv[0], kv[1]
         q_global = q_global.repeat(1, B_dim, 1, 1, 1)
         q = q_global.reshape(B_, self.num_heads, N, head_dim)
         q = q * self.scale
-        attn = (q @ k.transpose(-2, -1))
-        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
-            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)
+        attn = q @ k.transpose(-2, -1)
+        relative_position_bias = self.relative_position_bias_table[
+            self.relative_position_index.view(-1)
+        ].view(
+            self.window_size[0] * self.window_size[1],
+            self.window_size[0] * self.window_size[1],
+            -1,
+        )
         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
         attn = attn + relative_position_bias.unsqueeze(0)
         attn = self.softmax(attn)
@@ -139,6 +158,7 @@ class WindowAttentionGlobal(nn.Module):
         x = self.proj_drop(x)
         return x
 
+
 def _to_channel_last(x):
     """
     Args:
@@ -149,25 +169,30 @@ def _to_channel_last(x):
     """
     return x.permute(0, 2, 3, 1)
 
+
 def _to_channel_first(x):
     return x.permute(0, 3, 1, 2)
 
+
 def _to_query(x, N, num_heads, dim_head):
     B = x.shape[0]
     x = x.reshape(B, 1, N, num_heads, dim_head).permute(0, 1, 3, 2, 4)
     return x
 
+
 class Mlp(nn.Module):
     """
     Multi-Layer Perceptron (MLP) block
     """
 
-    def __init__(self,
-                 in_features,
-                 hidden_features=None,
-                 out_features=None,
-                 act_layer=nn.GELU,
-                 drop=0.):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
         """
         Args:
             in_features: input features dimension.
@@ -193,6 +218,7 @@ class Mlp(nn.Module):
         x = self.drop(x)
         return x
 
+
 def window_reverse(windows, window_size, H, W, h_w, w_w):
     """
     Args:
@@ -209,20 +235,40 @@ def window_reverse(windows, window_size, H, W, h_w, w_w):
     x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
     return x
 
+
 class feedforward(nn.Module):
-    def __init__(self, dim, window_size, mlp_ratio=4., act_layer=nn.GELU, drop=0., drop_path=0., layer_scale=None):
+    def __init__(
+        self,
+        dim,
+        window_size,
+        mlp_ratio=4.0,
+        act_layer=nn.GELU,
+        drop=0.0,
+        drop_path=0.0,
+        layer_scale=None,
+    ):
         super(feedforward, self).__init__()
         if layer_scale is not None and type(layer_scale) in [int, float]:
             self.layer_scale = True
-            self.gamma1 = nn.Parameter(layer_scale * torch.ones(dim), requires_grad=True)
-            self.gamma2 = nn.Parameter(layer_scale * torch.ones(dim), requires_grad=True)
+            self.gamma1 = nn.Parameter(
+                layer_scale * torch.ones(dim), requires_grad=True
+            )
+            self.gamma2 = nn.Parameter(
+                layer_scale * torch.ones(dim), requires_grad=True
+            )
         else:
             self.gamma1 = 1.0
             self.gamma2 = 1.0
         self.window_size = window_size
-        self.mlp = Mlp(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=drop,
+        )
         self.norm = nn.LayerNorm(dim)
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
     def forward(self, attn_windows, shortcut):
         B, H, W, C = shortcut.shape
         h_w = int(torch.div(H, self.window_size).item())
@@ -232,8 +278,17 @@ class feedforward(nn.Module):
         x = x + self.drop_path(self.gamma2 * self.mlp(self.norm(x)))
         return x
 
+
 class pyramid_trans_expr2(nn.Module):
-    def __init__(self, img_size=224, num_classes=8, window_size=[28,14,7], num_heads=[2, 4, 8], dims=[64, 128, 256], embed_dim=768):
+    def __init__(
+        self,
+        img_size=224,
+        num_classes=8,
+        window_size=[28, 14, 7],
+        num_heads=[2, 4, 8],
+        dims=[64, 128, 256],
+        embed_dim=768,
+    ):
         super().__init__()
 
         self.img_size = img_size
@@ -245,51 +300,99 @@ class pyramid_trans_expr2(nn.Module):
         self.window_size = window_size
         self.N = [win * win for win in window_size]
         self.face_landback = MobileFaceNet([112, 112], 136)
-        face_landback_checkpoint = torch.load(r'./pretrain/mobilefacenet_model_best.pth.tar',
-                                              map_location=lambda storage, loc: storage)
-        self.face_landback.load_state_dict(face_landback_checkpoint['state_dict'])
+        face_landback_checkpoint = torch.load(
+            r"./pretrain/mobilefacenet_model_best.pth.tar",
+            map_location=lambda storage, loc: storage,
+        )
+        self.face_landback.load_state_dict(face_landback_checkpoint["state_dict"])
 
         for param in self.face_landback.parameters():
             param.requires_grad = False
 
-        self.VIT = VisionTransformer(depth=2, embed_dim=embed_dim, num_classes=num_classes)
+        self.VIT = VisionTransformer(
+            depth=2, embed_dim=embed_dim, num_classes=num_classes
+        )
 
-        self.ir_back = Backbone(50, 0.0, 'ir')
-        ir_checkpoint = torch.load(r'./pretrain/ir50.pth', map_location=lambda storage, loc: storage)
+        self.ir_back = Backbone(50, 0.0, "ir")
+        ir_checkpoint = torch.load(
+            r"./pretrain/ir50.pth", map_location=lambda storage, loc: storage
+        )
 
         self.ir_back = load_pretrained_weights(self.ir_back, ir_checkpoint)
 
-        self.attn1 = WindowAttentionGlobal(dim=dims[0], num_heads=num_heads[0], window_size=window_size[0])
-        self.attn2 = WindowAttentionGlobal(dim=dims[1], num_heads=num_heads[1], window_size=window_size[1])
-        self.attn3 = WindowAttentionGlobal(dim=dims[2], num_heads=num_heads[2], window_size=window_size[2])
+        self.attn1 = WindowAttentionGlobal(
+            dim=dims[0], num_heads=num_heads[0], window_size=window_size[0]
+        )
+        self.attn2 = WindowAttentionGlobal(
+            dim=dims[1], num_heads=num_heads[1], window_size=window_size[1]
+        )
+        self.attn3 = WindowAttentionGlobal(
+            dim=dims[2], num_heads=num_heads[2], window_size=window_size[2]
+        )
         self.window1 = window(window_size=window_size[0], dim=dims[0])
         self.window2 = window(window_size=window_size[1], dim=dims[1])
         self.window3 = window(window_size=window_size[2], dim=dims[2])
-        self.conv1 = nn.Conv2d(in_channels=dims[0], out_channels=dims[0], kernel_size=3, stride=2, padding=1)
-        self.conv2 = nn.Conv2d(in_channels=dims[1], out_channels=dims[1], kernel_size=3, stride=2, padding=1)
-        self.conv3 = nn.Conv2d(in_channels=dims[2], out_channels=dims[2], kernel_size=3, stride=2, padding=1)
+        self.conv1 = nn.Conv2d(
+            in_channels=dims[0],
+            out_channels=dims[0],
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels=dims[1],
+            out_channels=dims[1],
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        )
+        self.conv3 = nn.Conv2d(
+            in_channels=dims[2],
+            out_channels=dims[2],
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        )
 
         dpr = [x.item() for x in torch.linspace(0, 0.5, 5)]
-        self.ffn1 = feedforward(dim=dims[0], window_size=window_size[0], layer_scale=1e-5, drop_path=dpr[0])
-        self.ffn2 = feedforward(dim=dims[1], window_size=window_size[1], layer_scale=1e-5, drop_path=dpr[1])
-        self.ffn3 = feedforward(dim=dims[2], window_size=window_size[2], layer_scale=1e-5, drop_path=dpr[2])
-
-        self.last_face_conv = nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3, padding=1)
-
-        self.embed_q = nn.Sequential(nn.Conv2d(dims[0], 768, kernel_size=3, stride=2, padding=1),
-                                     nn.Conv2d(768, 768, kernel_size=3, stride=2, padding=1))
-        self.embed_k = nn.Sequential(nn.Conv2d(dims[1], 768, kernel_size=3, stride=2, padding=1))
+        self.ffn1 = feedforward(
+            dim=dims[0], window_size=window_size[0], layer_scale=1e-5, drop_path=dpr[0]
+        )
+        self.ffn2 = feedforward(
+            dim=dims[1], window_size=window_size[1], layer_scale=1e-5, drop_path=dpr[1]
+        )
+        self.ffn3 = feedforward(
+            dim=dims[2], window_size=window_size[2], layer_scale=1e-5, drop_path=dpr[2]
+        )
+
+        self.last_face_conv = nn.Conv2d(
+            in_channels=512, out_channels=256, kernel_size=3, padding=1
+        )
+
+        self.embed_q = nn.Sequential(
+            nn.Conv2d(dims[0], 768, kernel_size=3, stride=2, padding=1),
+            nn.Conv2d(768, 768, kernel_size=3, stride=2, padding=1),
+        )
+        self.embed_k = nn.Sequential(
+            nn.Conv2d(dims[1], 768, kernel_size=3, stride=2, padding=1)
+        )
         self.embed_v = PatchEmbed(img_size=14, patch_size=14, in_c=256, embed_dim=768)
 
     def forward(self, x):
         x_face = F.interpolate(x, size=112)
-        x_face1 , x_face2, x_face3 = self.face_landback(x_face)
+        x_face1, x_face2, x_face3 = self.face_landback(x_face)
         x_face3 = self.last_face_conv(x_face3)
-        x_face1, x_face2, x_face3 = _to_channel_last(x_face1), _to_channel_last(x_face2), _to_channel_last(x_face3)
-
-        q1, q2, q3 = _to_query(x_face1, self.N[0], self.num_heads[0], self.dim_head[0]), \
-                     _to_query(x_face2, self.N[1], self.num_heads[1], self.dim_head[1]), \
-                     _to_query(x_face3, self.N[2], self.num_heads[2], self.dim_head[2])
+        x_face1, x_face2, x_face3 = (
+            _to_channel_last(x_face1),
+            _to_channel_last(x_face2),
+            _to_channel_last(x_face3),
+        )
+
+        q1, q2, q3 = (
+            _to_query(x_face1, self.N[0], self.num_heads[0], self.dim_head[0]),
+            _to_query(x_face2, self.N[1], self.num_heads[1], self.dim_head[1]),
+            _to_query(x_face3, self.N[2], self.num_heads[2], self.dim_head[2]),
+        )
 
         x_ir1, x_ir2, x_ir3 = self.ir_back(x)
         x_ir1, x_ir2, x_ir3 = self.conv1(x_ir1), self.conv2(x_ir2), self.conv3(x_ir3)
@@ -297,21 +400,34 @@ class pyramid_trans_expr2(nn.Module):
         x_window2, shortcut2 = self.window2(x_ir2)
         x_window3, shortcut3 = self.window3(x_ir3)
 
-        o1, o2, o3 = self.attn1(x_window1, q1), self.attn2(x_window2, q2), self.attn3(x_window3, q3)
+        o1, o2, o3 = (
+            self.attn1(x_window1, q1),
+            self.attn2(x_window2, q2),
+            self.attn3(x_window3, q3),
+        )
 
-        o1, o2, o3 = self.ffn1(o1, shortcut1), self.ffn2(o2, shortcut2), self.ffn3(o3, shortcut3)
+        o1, o2, o3 = (
+            self.ffn1(o1, shortcut1),
+            self.ffn2(o2, shortcut2),
+            self.ffn3(o3, shortcut3),
+        )
 
         o1, o2, o3 = _to_channel_first(o1), _to_channel_first(o2), _to_channel_first(o3)
 
-        o1, o2, o3 = self.embed_q(o1).flatten(2).transpose(1, 2), self.embed_k(o2).flatten(2).transpose(1, 2), self.embed_v(o3)
+        o1, o2, o3 = (
+            self.embed_q(o1).flatten(2).transpose(1, 2),
+            self.embed_k(o2).flatten(2).transpose(1, 2),
+            self.embed_v(o3),
+        )
 
         o = torch.cat([o1, o2, o3], dim=1)
 
         out = self.VIT(o)
         return out
 
+
 def compute_param_flop():
     model = pyramid_trans_expr2()
-    img = torch.rand(size=(1,3,224,224))
+    img = torch.rand(size=(1, 3, 224, 224))
     flops, params = profile(model, inputs=(img,))
-    print(f'flops:{flops/1000**3}G,params:{params/1000**2}M')
+    print(f"flops:{flops/1000**3}G,params:{params/1000**2}M")

FER/models/vit_model.py

@@ -2,6 +2,7 @@
 original code from rwightman:
 https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
 """
+
 from functools import partial
 from collections import OrderedDict
 
@@ -23,16 +24,24 @@ import torch.hub
 from functools import partial
 import math
 
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-from timm.models.registry import register_model
+from timm.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models import register_model
 from timm.models.vision_transformer import _cfg, Mlp, Block
 # from .ir50 import Backbone
 
 
 def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
     """3x3 convolution with padding"""
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
 
 
 def conv1x1(in_planes, out_planes, stride=1):
@@ -40,7 +49,7 @@ def conv1x1(in_planes, out_planes, stride=1):
     return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
 
 
-def drop_path(x, drop_prob: float = 0., training: bool = False):
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
     """
     Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
     This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
@@ -49,10 +58,12 @@ def drop_path(x, drop_prob: float = 0., training: bool = False):
     changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
     'survival rate' as the argument.
     """
-    if drop_prob == 0. or not training:
+    if drop_prob == 0.0 or not training:
         return x
     keep_prob = 1 - drop_prob
-    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
     random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
     random_tensor.floor_()  # binarize
     output = x.div(keep_prob) * random_tensor
@@ -60,7 +71,7 @@ def drop_path(x, drop_prob: float = 0., training: bool = False):
 
 
 class BasicBlock(nn.Module):
-    __constants__ = ['downsample']
+    __constants__ = ["downsample"]
 
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super(BasicBlock, self).__init__()
@@ -109,7 +120,9 @@ class PatchEmbed(nn.Module):
     2D Image to Patch Embedding
     """
 
-    def __init__(self, img_size=14, patch_size=16, in_c=256, embed_dim=768, norm_layer=None):
+    def __init__(
+        self, img_size=14, patch_size=16, in_c=256, embed_dim=768, norm_layer=None
+    ):
         super().__init__()
         img_size = (img_size, img_size)
         patch_size = (patch_size, patch_size)
@@ -135,29 +148,36 @@ class PatchEmbed(nn.Module):
 
 
 class Attention(nn.Module):
-    def __init__(self,
-                 dim, in_chans,  # 输入token的dim
-                 num_heads=8,
-                 qkv_bias=False,
-                 qk_scale=None,
-                 attn_drop_ratio=0.,
-                 proj_drop_ratio=0.):
+    def __init__(
+        self,
+        dim,
+        in_chans,  # 输入token的dim
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop_ratio=0.0,
+        proj_drop_ratio=0.0,
+    ):
         super(Attention, self).__init__()
         self.num_heads = 8
         self.img_chanel = in_chans + 1
         head_dim = dim // num_heads
-        self.scale = head_dim ** -0.5
+        self.scale = head_dim**-0.5
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
         self.attn_drop = nn.Dropout(attn_drop_ratio)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop_ratio)
 
     def forward(self, x):
-        x_img = x[:, :self.img_chanel, :]
+        x_img = x[:, : self.img_chanel, :]
         # [batch_size, num_patches + 1, total_embed_dim]
         B, N, C = x_img.shape
         # print(C)
-        qkv = self.qkv(x_img).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        qkv = (
+            self.qkv(x_img)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
         q, k, v = qkv[0], qkv[1], qkv[2]
         # k, v = kv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
         # q = x_img.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
@@ -193,7 +213,7 @@ class Attention(nn.Module):
 
 
 class AttentionBlock(nn.Module):
-    __constants__ = ['downsample']
+    __constants__ = ["downsample"]
 
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super(AttentionBlock, self).__init__()
@@ -234,7 +254,14 @@ class Mlp(nn.Module):
     MLP as used in Vision Transformer, MLP-Mixer and related networks
     """
 
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
@@ -253,29 +280,46 @@ class Mlp(nn.Module):
 
 
 class Block(nn.Module):
-    def __init__(self,
-                 dim, in_chans,
-                 num_heads,
-                 mlp_ratio=4.,
-                 qkv_bias=False,
-                 qk_scale=None,
-                 drop_ratio=0.,
-                 attn_drop_ratio=0.,
-                 drop_path_ratio=0.,
-                 act_layer=nn.GELU,
-                 norm_layer=nn.LayerNorm):
+    def __init__(
+        self,
+        dim,
+        in_chans,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_ratio=0.0,
+        attn_drop_ratio=0.0,
+        drop_path_ratio=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
         super(Block, self).__init__()
         self.norm1 = norm_layer(dim)
         self.img_chanel = in_chans + 1
 
         self.conv = nn.Conv1d(self.img_chanel, self.img_chanel, 1)
-        self.attn = Attention(dim, in_chans=in_chans, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
-                              attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
+        self.attn = Attention(
+            dim,
+            in_chans=in_chans,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop_ratio=attn_drop_ratio,
+            proj_drop_ratio=drop_ratio,
+        )
         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
-        self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
+        self.drop_path = (
+            DropPath(drop_path_ratio) if drop_path_ratio > 0.0 else nn.Identity()
+        )
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop_ratio,
+        )
 
     def forward(self, x):
         # x = x + self.drop_path(self.attn(self.norm1(x)))
@@ -308,8 +352,9 @@ class ClassificationHead(nn.Module):
 
 def load_pretrained_weights(model, checkpoint):
     import collections
-    if 'state_dict' in checkpoint:
-        state_dict = checkpoint['state_dict']
+
+    if "state_dict" in checkpoint:
+        state_dict = checkpoint["state_dict"]
     else:
         state_dict = checkpoint
     model_dict = model.state_dict()
@@ -318,7 +363,7 @@ def load_pretrained_weights(model, checkpoint):
     for k, v in state_dict.items():
         # If the pretrained state_dict was saved as nn.DataParallel,
         # keys would contain "module.", which should be ignored.
-        if k.startswith('module.'):
+        if k.startswith("module."):
             k = k[7:]
         if k in model_dict and model_dict[k].size() == v.size():
             new_state_dict[k] = v
@@ -329,9 +374,10 @@ def load_pretrained_weights(model, checkpoint):
     model_dict.update(new_state_dict)
 
     model.load_state_dict(model_dict)
-    print('load_weight', len(matched_layers))
+    print("load_weight", len(matched_layers))
     return model
 
+
 class eca_block(nn.Module):
     def __init__(self, channel=128, b=1, gamma=2):
         super(eca_block, self).__init__()
@@ -339,7 +385,9 @@ class eca_block(nn.Module):
         kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
 
         self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
+        self.conv = nn.Conv1d(
+            1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False
+        )
         self.sigmoid = nn.Sigmoid()
 
     def forward(self, x):
@@ -347,6 +395,8 @@ class eca_block(nn.Module):
         y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
         y = self.sigmoid(y)
         return x * y.expand_as(x)
+
+
 #
 #
 # class IR20(nn.Module):
@@ -484,7 +534,9 @@ class eca_block(nn.Module):
         kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
 
         self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
+        self.conv = nn.Conv1d(
+            1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False
+        )
         self.sigmoid = nn.Sigmoid()
 
     def forward(self, x):
@@ -493,6 +545,7 @@ class eca_block(nn.Module):
         y = self.sigmoid(y)
         return x * y.expand_as(x)
 
+
 class SE_block(nn.Module):
     def __init__(self, input_dim: int):
         super().__init__()
@@ -511,11 +564,27 @@ class SE_block(nn.Module):
 
 
 class VisionTransformer(nn.Module):
-    def __init__(self, img_size=14, patch_size=14, in_c=147, num_classes=7,
-                 embed_dim=768, depth=6, num_heads=8, mlp_ratio=4.0, qkv_bias=True,
-                 qk_scale=None, representation_size=None, distilled=False, drop_ratio=0.,
-                 attn_drop_ratio=0., drop_path_ratio=0., embed_layer=PatchEmbed, norm_layer=None,
-                 act_layer=None):
+    def __init__(
+        self,
+        img_size=14,
+        patch_size=14,
+        in_c=147,
+        num_classes=7,
+        embed_dim=768,
+        depth=6,
+        num_heads=8,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        representation_size=None,
+        distilled=False,
+        drop_ratio=0.0,
+        attn_drop_ratio=0.0,
+        drop_path_ratio=0.0,
+        embed_layer=PatchEmbed,
+        norm_layer=None,
+        act_layer=None,
+    ):
         """
         Args:
             img_size (int, tuple): input image size
@@ -538,7 +607,9 @@ class VisionTransformer(nn.Module):
         """
         super(VisionTransformer, self).__init__()
         self.num_classes = num_classes
-        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_features = self.embed_dim = (
+            embed_dim  # num_features for consistency with other models
+        )
         self.num_tokens = 2 if distilled else 1
         norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
         act_layer = act_layer or nn.GELU
@@ -549,18 +620,20 @@ class VisionTransformer(nn.Module):
 
         self.se_block = SE_block(input_dim=embed_dim)
 
-
-        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_c=256, embed_dim=768)
+        self.patch_embed = embed_layer(
+            img_size=img_size, patch_size=patch_size, in_c=256, embed_dim=768
+        )
         num_patches = self.patch_embed.num_patches
         self.head = ClassificationHead(input_dim=embed_dim, target_dim=self.num_classes)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
+        self.dist_token = (
+            nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
+        )
         # self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
         self.pos_drop = nn.Dropout(p=drop_ratio)
         # self.IR = IR()
         self.eca_block = eca_block()
 
-
         # self.ir_back = Backbone(50, 0.0, 'ir')
         # ir_checkpoint = torch.load('./models/pretrain/ir50.pth', map_location=lambda storage, loc: storage)
         # # ir_checkpoint = ir_checkpoint["model"]
@@ -570,24 +643,41 @@ class VisionTransformer(nn.Module):
         self.IRLinear1 = nn.Linear(1024, 768)
         self.IRLinear2 = nn.Linear(768, 512)
         self.eca_block = eca_block()
-        dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)]  # stochastic depth decay rule
-        self.blocks = nn.Sequential(*[
-            Block(dim=embed_dim, in_chans=in_c, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
-                  qk_scale=qk_scale,
-                  drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],
-                  norm_layer=norm_layer, act_layer=act_layer)
-            for i in range(depth)
-        ])
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_ratio, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(
+            *[
+                Block(
+                    dim=embed_dim,
+                    in_chans=in_c,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop_ratio=drop_ratio,
+                    attn_drop_ratio=attn_drop_ratio,
+                    drop_path_ratio=dpr[i],
+                    norm_layer=norm_layer,
+                    act_layer=act_layer,
+                )
+                for i in range(depth)
+            ]
+        )
         self.norm = norm_layer(embed_dim)
 
         # Representation layer
         if representation_size and not distilled:
             self.has_logits = True
             self.num_features = representation_size
-            self.pre_logits = nn.Sequential(OrderedDict([
-                ("fc", nn.Linear(embed_dim, representation_size)),
-                ("act", nn.Tanh())
-            ]))
+            self.pre_logits = nn.Sequential(
+                OrderedDict(
+                    [
+                        ("fc", nn.Linear(embed_dim, representation_size)),
+                        ("act", nn.Tanh()),
+                    ]
+                )
+            )
         else:
             self.has_logits = False
             self.pre_logits = nn.Identity()
@@ -596,7 +686,11 @@ class VisionTransformer(nn.Module):
         # self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
         self.head_dist = None
         if distilled:
-            self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()
+            self.head_dist = (
+                nn.Linear(self.embed_dim, self.num_classes)
+                if num_classes > 0
+                else nn.Identity()
+            )
 
         # Weight init
         nn.init.trunc_normal_(self.pos_embed, std=0.02)
@@ -616,7 +710,9 @@ class VisionTransformer(nn.Module):
         if self.dist_token is None:
             x = torch.cat((cls_token, x), dim=1)  # [B, 197, 768]
         else:
-            x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
+            x = torch.cat(
+                (cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1
+            )
         # print(x.shape)
         x = self.pos_drop(x + self.pos_embed)
         x = self.blocks(x)
@@ -627,7 +723,6 @@ class VisionTransformer(nn.Module):
             return x[:, 0], x[:, 1]
 
     def forward(self, x):
-
         # B = x.shape[0]
         # print(x)
         # x = self.eca_block(x)
@@ -680,7 +775,7 @@ def _init_vit_weights(m):
     :param m: module
     """
     if isinstance(m, nn.Linear):
-        nn.init.trunc_normal_(m.weight, std=.01)
+        nn.init.trunc_normal_(m.weight, std=0.01)
         if m.bias is not None:
             nn.init.zeros_(m.bias)
     elif isinstance(m, nn.Conv2d):
@@ -699,13 +794,15 @@ def vit_base_patch16_224(num_classes: int = 7):
     weights ported from official Google JAX impl:
     链接: https://pan.baidu.com/s/1zqb08naP0RPqqfSXfkB2EA  密码: eu9f
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=None,
+        num_classes=num_classes,
+    )
 
     return model
 
@@ -717,13 +814,15 @@ def vit_base_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch16_224_in21k-e5005f0a.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=768 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=768 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -734,13 +833,15 @@ def vit_base_patch32_224(num_classes: int = 1000):
     weights ported from official Google JAX impl:
     链接: https://pan.baidu.com/s/1hCv0U8pQomwAtHBYc4hmZg  密码: s5hl
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=32,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=32,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -751,13 +852,15 @@ def vit_base_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch32_224_in21k-8db57226.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=32,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=768 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=32,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=768 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -768,13 +871,15 @@ def vit_large_patch16_224(num_classes: int = 1000):
     weights ported from official Google JAX impl:
     链接: https://pan.baidu.com/s/1cxBgZJJ6qUWPSBNcE4TdRQ  密码: qqt8
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=1024,
-                              depth=24,
-                              num_heads=16,
-                              representation_size=None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        representation_size=None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -785,13 +890,15 @@ def vit_large_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = Tru
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch16_224_in21k-606da67d.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=1024,
-                              depth=24,
-                              num_heads=16,
-                              representation_size=1024 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        representation_size=1024 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -802,13 +909,15 @@ def vit_large_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = Tru
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=32,
-                              embed_dim=1024,
-                              depth=24,
-                              num_heads=16,
-                              representation_size=1024 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=32,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        representation_size=1024 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -818,11 +927,13 @@ def vit_huge_patch14_224_in21k(num_classes: int = 21843, has_logits: bool = True
     ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
     NOTE: converted weights not currently available, too large for github release hosting.
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=14,
-                              embed_dim=1280,
-                              depth=32,
-                              num_heads=16,
-                              representation_size=1280 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=14,
+        embed_dim=1280,
+        depth=32,
+        num_heads=16,
+        representation_size=1280 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model

FER/models/vit_model_8.py

@@ -2,6 +2,7 @@
 original code from rwightman:
 https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
 """
+
 from functools import partial
 from collections import OrderedDict
 
@@ -23,16 +24,24 @@ import torch.hub
 from functools import partial
 import math
 
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-from timm.models.registry import register_model
+from timm.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models import register_model
 from timm.models.vision_transformer import _cfg, Mlp, Block
 from .ir50 import Backbone
 
 
 def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
     """3x3 convolution with padding"""
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
 
 
 def conv1x1(in_planes, out_planes, stride=1):
@@ -40,7 +49,7 @@ def conv1x1(in_planes, out_planes, stride=1):
     return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
 
 
-def drop_path(x, drop_prob: float = 0., training: bool = False):
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
     """
     Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
     This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
@@ -49,10 +58,12 @@ def drop_path(x, drop_prob: float = 0., training: bool = False):
     changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
     'survival rate' as the argument.
     """
-    if drop_prob == 0. or not training:
+    if drop_prob == 0.0 or not training:
         return x
     keep_prob = 1 - drop_prob
-    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
     random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
     random_tensor.floor_()  # binarize
     output = x.div(keep_prob) * random_tensor
@@ -60,7 +71,7 @@ def drop_path(x, drop_prob: float = 0., training: bool = False):
 
 
 class BasicBlock(nn.Module):
-    __constants__ = ['downsample']
+    __constants__ = ["downsample"]
 
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super(BasicBlock, self).__init__()
@@ -109,7 +120,9 @@ class PatchEmbed(nn.Module):
     2D Image to Patch Embedding
     """
 
-    def __init__(self, img_size=14, patch_size=16, in_c=256, embed_dim=768, norm_layer=None):
+    def __init__(
+        self, img_size=14, patch_size=16, in_c=256, embed_dim=768, norm_layer=None
+    ):
         super().__init__()
         img_size = (img_size, img_size)
         patch_size = (patch_size, patch_size)
@@ -135,29 +148,36 @@ class PatchEmbed(nn.Module):
 
 
 class Attention(nn.Module):
-    def __init__(self,
-                 dim, in_chans,  # 输入token的dim
-                 num_heads=8,
-                 qkv_bias=False,
-                 qk_scale=None,
-                 attn_drop_ratio=0.,
-                 proj_drop_ratio=0.):
+    def __init__(
+        self,
+        dim,
+        in_chans,  # 输入token的dim
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop_ratio=0.0,
+        proj_drop_ratio=0.0,
+    ):
         super(Attention, self).__init__()
         self.num_heads = 8
         self.img_chanel = in_chans + 1
         head_dim = dim // num_heads
-        self.scale = head_dim ** -0.5
+        self.scale = head_dim**-0.5
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
         self.attn_drop = nn.Dropout(attn_drop_ratio)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop_ratio)
 
     def forward(self, x):
-        x_img = x[:, :self.img_chanel, :]
+        x_img = x[:, : self.img_chanel, :]
         # [batch_size, num_patches + 1, total_embed_dim]
         B, N, C = x_img.shape
         # print(C)
-        qkv = self.qkv(x_img).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        qkv = (
+            self.qkv(x_img)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
         q, k, v = qkv[0], qkv[1], qkv[2]
         # k, v = kv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
         # q = x_img.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
@@ -193,7 +213,7 @@ class Attention(nn.Module):
 
 
 class AttentionBlock(nn.Module):
-    __constants__ = ['downsample']
+    __constants__ = ["downsample"]
 
     def __init__(self, inplanes, planes, stride=1, downsample=None):
         super(AttentionBlock, self).__init__()
@@ -234,7 +254,14 @@ class Mlp(nn.Module):
     MLP as used in Vision Transformer, MLP-Mixer and related networks
     """
 
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
@@ -253,29 +280,46 @@ class Mlp(nn.Module):
 
 
 class Block(nn.Module):
-    def __init__(self,
-                 dim, in_chans,
-                 num_heads,
-                 mlp_ratio=4.,
-                 qkv_bias=False,
-                 qk_scale=None,
-                 drop_ratio=0.,
-                 attn_drop_ratio=0.,
-                 drop_path_ratio=0.,
-                 act_layer=nn.GELU,
-                 norm_layer=nn.LayerNorm):
+    def __init__(
+        self,
+        dim,
+        in_chans,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_ratio=0.0,
+        attn_drop_ratio=0.0,
+        drop_path_ratio=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
         super(Block, self).__init__()
         self.norm1 = norm_layer(dim)
         self.img_chanel = in_chans + 1
 
         self.conv = nn.Conv1d(self.img_chanel, self.img_chanel, 1)
-        self.attn = Attention(dim, in_chans=in_chans, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
-                              attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
+        self.attn = Attention(
+            dim,
+            in_chans=in_chans,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop_ratio=attn_drop_ratio,
+            proj_drop_ratio=drop_ratio,
+        )
         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
-        self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
+        self.drop_path = (
+            DropPath(drop_path_ratio) if drop_path_ratio > 0.0 else nn.Identity()
+        )
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop_ratio,
+        )
 
     def forward(self, x):
         # x = x + self.drop_path(self.attn(self.norm1(x)))
@@ -308,8 +352,9 @@ class ClassificationHead(nn.Module):
 
 def load_pretrained_weights(model, checkpoint):
     import collections
-    if 'state_dict' in checkpoint:
-        state_dict = checkpoint['state_dict']
+
+    if "state_dict" in checkpoint:
+        state_dict = checkpoint["state_dict"]
     else:
         state_dict = checkpoint
     model_dict = model.state_dict()
@@ -318,7 +363,7 @@ def load_pretrained_weights(model, checkpoint):
     for k, v in state_dict.items():
         # If the pretrained state_dict was saved as nn.DataParallel,
         # keys would contain "module.", which should be ignored.
-        if k.startswith('module.'):
+        if k.startswith("module."):
             k = k[7:]
         if k in model_dict and model_dict[k].size() == v.size():
             new_state_dict[k] = v
@@ -329,9 +374,10 @@ def load_pretrained_weights(model, checkpoint):
     model_dict.update(new_state_dict)
 
     model.load_state_dict(model_dict)
-    print('load_weight', len(matched_layers))
+    print("load_weight", len(matched_layers))
     return model
 
+
 class eca_block(nn.Module):
     def __init__(self, channel=128, b=1, gamma=2):
         super(eca_block, self).__init__()
@@ -339,7 +385,9 @@ class eca_block(nn.Module):
         kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
 
         self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
+        self.conv = nn.Conv1d(
+            1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False
+        )
         self.sigmoid = nn.Sigmoid()
 
     def forward(self, x):
@@ -347,6 +395,8 @@ class eca_block(nn.Module):
         y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
         y = self.sigmoid(y)
         return x * y.expand_as(x)
+
+
 #
 #
 # class IR20(nn.Module):
@@ -484,7 +534,9 @@ class eca_block(nn.Module):
         kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
 
         self.avg_pool = nn.AdaptiveAvgPool2d(1)
-        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
+        self.conv = nn.Conv1d(
+            1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False
+        )
         self.sigmoid = nn.Sigmoid()
 
     def forward(self, x):
@@ -493,6 +545,7 @@ class eca_block(nn.Module):
         y = self.sigmoid(y)
         return x * y.expand_as(x)
 
+
 class SE_block(nn.Module):
     def __init__(self, input_dim: int):
         super().__init__()
@@ -511,11 +564,27 @@ class SE_block(nn.Module):
 
 
 class VisionTransformer(nn.Module):
-    def __init__(self, img_size=14, patch_size=14, in_c=147, num_classes=8,
-                 embed_dim=768, depth=6, num_heads=8, mlp_ratio=4.0, qkv_bias=True,
-                 qk_scale=None, representation_size=None, distilled=False, drop_ratio=0.,
-                 attn_drop_ratio=0., drop_path_ratio=0., embed_layer=PatchEmbed, norm_layer=None,
-                 act_layer=None):
+    def __init__(
+        self,
+        img_size=14,
+        patch_size=14,
+        in_c=147,
+        num_classes=8,
+        embed_dim=768,
+        depth=6,
+        num_heads=8,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        representation_size=None,
+        distilled=False,
+        drop_ratio=0.0,
+        attn_drop_ratio=0.0,
+        drop_path_ratio=0.0,
+        embed_layer=PatchEmbed,
+        norm_layer=None,
+        act_layer=None,
+    ):
         """
         Args:
             img_size (int, tuple): input image size
@@ -538,7 +607,9 @@ class VisionTransformer(nn.Module):
         """
         super(VisionTransformer, self).__init__()
         self.num_classes = num_classes
-        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_features = self.embed_dim = (
+            embed_dim  # num_features for consistency with other models
+        )
         self.num_tokens = 2 if distilled else 1
         norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
         act_layer = act_layer or nn.GELU
@@ -549,18 +620,20 @@ class VisionTransformer(nn.Module):
 
         self.se_block = SE_block(input_dim=embed_dim)
 
-
-        self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_c=256, embed_dim=768)
+        self.patch_embed = embed_layer(
+            img_size=img_size, patch_size=patch_size, in_c=256, embed_dim=768
+        )
         num_patches = self.patch_embed.num_patches
         self.head = ClassificationHead(input_dim=embed_dim, target_dim=self.num_classes)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
+        self.dist_token = (
+            nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
+        )
         # self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
         self.pos_drop = nn.Dropout(p=drop_ratio)
         # self.IR = IR()
         self.eca_block = eca_block()
 
-
         # self.ir_back = Backbone(50, 0.0, 'ir')
         # ir_checkpoint = torch.load('./models/pretrain/ir50.pth', map_location=lambda storage, loc: storage)
         # # ir_checkpoint = ir_checkpoint["model"]
@@ -570,24 +643,41 @@ class VisionTransformer(nn.Module):
         self.IRLinear1 = nn.Linear(1024, 768)
         self.IRLinear2 = nn.Linear(768, 512)
         self.eca_block = eca_block()
-        dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)]  # stochastic depth decay rule
-        self.blocks = nn.Sequential(*[
-            Block(dim=embed_dim, in_chans=in_c, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
-                  qk_scale=qk_scale,
-                  drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],
-                  norm_layer=norm_layer, act_layer=act_layer)
-            for i in range(depth)
-        ])
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_ratio, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(
+            *[
+                Block(
+                    dim=embed_dim,
+                    in_chans=in_c,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop_ratio=drop_ratio,
+                    attn_drop_ratio=attn_drop_ratio,
+                    drop_path_ratio=dpr[i],
+                    norm_layer=norm_layer,
+                    act_layer=act_layer,
+                )
+                for i in range(depth)
+            ]
+        )
         self.norm = norm_layer(embed_dim)
 
         # Representation layer
         if representation_size and not distilled:
             self.has_logits = True
             self.num_features = representation_size
-            self.pre_logits = nn.Sequential(OrderedDict([
-                ("fc", nn.Linear(embed_dim, representation_size)),
-                ("act", nn.Tanh())
-            ]))
+            self.pre_logits = nn.Sequential(
+                OrderedDict(
+                    [
+                        ("fc", nn.Linear(embed_dim, representation_size)),
+                        ("act", nn.Tanh()),
+                    ]
+                )
+            )
         else:
             self.has_logits = False
             self.pre_logits = nn.Identity()
@@ -596,7 +686,11 @@ class VisionTransformer(nn.Module):
         # self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
         self.head_dist = None
         if distilled:
-            self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()
+            self.head_dist = (
+                nn.Linear(self.embed_dim, self.num_classes)
+                if num_classes > 0
+                else nn.Identity()
+            )
 
         # Weight init
         nn.init.trunc_normal_(self.pos_embed, std=0.02)
@@ -616,7 +710,9 @@ class VisionTransformer(nn.Module):
         if self.dist_token is None:
             x = torch.cat((cls_token, x), dim=1)  # [B, 197, 768]
         else:
-            x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
+            x = torch.cat(
+                (cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1
+            )
         # print(x.shape)
         x = self.pos_drop(x + self.pos_embed)
         x = self.blocks(x)
@@ -627,7 +723,6 @@ class VisionTransformer(nn.Module):
             return x[:, 0], x[:, 1]
 
     def forward(self, x):
-
         # B = x.shape[0]
         # print(x)
         # x = self.eca_block(x)
@@ -680,7 +775,7 @@ def _init_vit_weights(m):
     :param m: module
     """
     if isinstance(m, nn.Linear):
-        nn.init.trunc_normal_(m.weight, std=.01)
+        nn.init.trunc_normal_(m.weight, std=0.01)
         if m.bias is not None:
             nn.init.zeros_(m.bias)
     elif isinstance(m, nn.Conv2d):
@@ -699,13 +794,15 @@ def vit_base_patch16_224(num_classes: int = 7):
     weights ported from official Google JAX impl:
     链接: https://pan.baidu.com/s/1zqb08naP0RPqqfSXfkB2EA  密码: eu9f
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=None,
+        num_classes=num_classes,
+    )
 
     return model
 
@@ -717,13 +814,15 @@ def vit_base_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch16_224_in21k-e5005f0a.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=768 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=768 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -734,13 +833,15 @@ def vit_base_patch32_224(num_classes: int = 1000):
     weights ported from official Google JAX impl:
     链接: https://pan.baidu.com/s/1hCv0U8pQomwAtHBYc4hmZg  密码: s5hl
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=32,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=32,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -751,13 +852,15 @@ def vit_base_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch32_224_in21k-8db57226.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=32,
-                              embed_dim=768,
-                              depth=12,
-                              num_heads=12,
-                              representation_size=768 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=32,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        representation_size=768 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -768,13 +871,15 @@ def vit_large_patch16_224(num_classes: int = 1000):
     weights ported from official Google JAX impl:
     链接: https://pan.baidu.com/s/1cxBgZJJ6qUWPSBNcE4TdRQ  密码: qqt8
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=1024,
-                              depth=24,
-                              num_heads=16,
-                              representation_size=None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        representation_size=None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -785,13 +890,15 @@ def vit_large_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = Tru
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch16_224_in21k-606da67d.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=16,
-                              embed_dim=1024,
-                              depth=24,
-                              num_heads=16,
-                              representation_size=1024 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=16,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        representation_size=1024 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -802,13 +909,15 @@ def vit_large_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = Tru
     weights ported from official Google JAX impl:
     https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=32,
-                              embed_dim=1024,
-                              depth=24,
-                              num_heads=16,
-                              representation_size=1024 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=32,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        representation_size=1024 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model
 
 
@@ -818,11 +927,13 @@ def vit_huge_patch14_224_in21k(num_classes: int = 21843, has_logits: bool = True
     ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.
     NOTE: converted weights not currently available, too large for github release hosting.
     """
-    model = VisionTransformer(img_size=224,
-                              patch_size=14,
-                              embed_dim=1280,
-                              depth=32,
-                              num_heads=16,
-                              representation_size=1280 if has_logits else None,
-                              num_classes=num_classes)
+    model = VisionTransformer(
+        img_size=224,
+        patch_size=14,
+        embed_dim=1280,
+        depth=32,
+        num_heads=16,
+        representation_size=1280 if has_logits else None,
+        num_classes=num_classes,
+    )
     return model

FER/prediction.py

@@ -48,7 +48,7 @@ def main():
                 )
             )
         else:
-            print("=> no checkpoint found at '{}'".format(model_path))
+            print("[!] prediction.py => no checkpoint found at '{}'".format(model_path))
         predict(model, image_path=image_arr)
         return