def train(conf): """Total training procedure. """ conf.device = torch.device('cuda:0') criterion = torch.nn.CrossEntropyLoss().cuda(conf.device) backbone_factory = BackboneFactory(conf.backbone_type, conf.backbone_conf_file) probe_net = backbone_factory.get_backbone() gallery_net = backbone_factory.get_backbone() head_factory = HeadFactory(conf.head_type, conf.head_conf_file) prototype = head_factory.get_head().cuda(conf.device) probe_net = torch.nn.DataParallel(probe_net).cuda() gallery_net = torch.nn.DataParallel(gallery_net).cuda() optimizer = optim.SGD(probe_net.parameters(), lr=conf.lr, momentum=conf.momentum, weight_decay=5e-4) lr_schedule = optim.lr_scheduler.MultiStepLR(optimizer, milestones=conf.milestones, gamma=0.1) if conf.resume: probe_net.load_state_dict(torch.load(args.pretrain_model)) moving_average(probe_net, gallery_net, 0) probe_net.train() gallery_net.eval().apply(train_BN) exclude_id_set = set() loss_meter = AverageMeter() for epoch in range(conf.epoches): data_loader = DataLoader( ImageDataset_SST(conf.data_root, conf.train_file, exclude_id_set), conf.batch_size, True, num_workers = 4, drop_last = True) exclude_id_set = train_one_epoch(data_loader, probe_net, gallery_net, prototype, optimizer, criterion, epoch, conf, loss_meter) lr_schedule.step()
def train(conf): """Total training procedure. """ data_loader = DataLoader(ImageDataset_SST(conf.data_root, conf.train_file), conf.batch_size, True, num_workers=4) conf.device = torch.device('cuda:0') #criterion = OnlineContrastiveLoss(margin=2.5, pair_selector=HardNegativePairSelector(cpu=False)).cuda(torch.device('cuda:0')) triplet_selector = FunctionNegativeTripletSelector( margin=2.5, negative_selection_fn=random_hard_negative, cpu=False) criterion = OnlineTripletLoss(margin=2.5, triplet_selector=triplet_selector).cuda( conf.device) backbone_factory = BackboneFactory(conf.backbone_type, conf.backbone_conf_file) model = backbone_factory.get_backbone() if conf.resume: model.load_state_dict(torch.load(args.pretrain_model)) model = torch.nn.DataParallel(model).cuda() parameters = [p for p in model.parameters() if p.requires_grad] optimizer = optim.SGD(parameters, lr=conf.lr, momentum=conf.momentum, weight_decay=1e-4) lr_schedule = optim.lr_scheduler.MultiStepLR(optimizer, milestones=conf.milestones, gamma=0.1) loss_meter = AverageMeter() model.train() for epoch in range(conf.epoches): train_one_epoch(data_loader, model, optimizer, criterion, epoch, loss_meter, conf) lr_schedule.step() if conf.evaluate: conf.evaluator.evaluate(model)
""" @author: Jun Wang @date: 20201012 @contact: [email protected] """ import sys import torch from thop import profile from thop import clever_format sys.path.append('..') from backbone.backbone_def import BackboneFactory backbone_type = 'MobileFaceNet' #backbone_type = 'ResNet' #backbone_type = 'EfficientNet' #backbone_type = 'HRNet' #backbone_type = 'GhostNet' #backbone_type = 'AttentionNet' backbone_conf_file = '../training_mode/backbone_conf.yaml' backbone_factory = BackboneFactory(backbone_type, backbone_conf_file) backbone = backbone_factory.get_backbone() input = torch.randn(1, 3, 112, 112) macs, params = profile(backbone, inputs=(input, )) macs, params = clever_format([macs, params], "%.2f") print('backbone type: ', backbone_type) print('Params: ', params) print('Macs: ', macs)