lichen/ieemoo-ai-contrast
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

多机并行计算

Browse Source
This commit is contained in:
lee
2025-08-14 10:09:54 +08:00
parent bc896fc688
commit 99a204ee22
18 changed files with 105 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
configs/utils.py
View File

@ -1,4 +1,4 @@
from model import (resnet18, resnet34, resnet50, mobilevit_s, MobileNetV3_Small, MobileNetV3_Large, mobilenet_v1,
from model import (resnet18, resnet34, resnet50, resnet101, mobilevit_s, MobileNetV3_Small, MobileNetV3_Large, mobilenet_v1,
PPLCNET_x1_0, PPLCNET_x0_5, PPLCNET_x2_5)
from timm.models import vit_base_patch16_224 as vit_base_16
from model.metric import ArcFace, CosFace
@ -13,9 +13,10 @@ class trainer_tools:
def get_backbone(self):
backbone_mapping = {
'resnet18': lambda: resnet18(scale=self.conf['models']['channel_ratio']),
'resnet34': lambda: resnet34(scale=self.conf['models']['channel_ratio']),
'resnet50': lambda: resnet50(scale=self.conf['models']['channel_ratio']),
'resnet18': lambda: resnet18(scale=self.conf['models']['channel_ratio'], pretrained=True),
'resnet34': lambda: resnet34(scale=self.conf['models']['channel_ratio'], pretrained=True),
'resnet50': lambda: resnet50(scale=self.conf['models']['channel_ratio'], pretrained=True),
'resnet101': lambda: resnet101(scale=self.conf['models']['channel_ratio'], pretrained=True),
'mobilevit_s': lambda: mobilevit_s(),
'mobilenetv3_small': lambda: MobileNetV3_Small(),
'PPLCNET_x1_0': lambda: PPLCNET_x1_0(),