import torch from model.efficientnet.model import EfficientNet net = EfficientNet.from_name('efficientnet-b4', override_params={'num_classes': 2}) model_dict = net.state_dict() weight = {} for key, value in torch.load( '/root/data/wzy/lxrlocker/scale2models_lxr/PAI_tracking_45part_0.1resize_epoch_55_checkpoint.pth' ).items(): key = key.replace('backbone.', '') if '_fc' in key: continue key = key.replace('last_linear', '_fc') weight[key] = value model_dict.update(weight) net.load_state_dict(model_dict) torch.save(net.state_dict(), "/root/data/wzy/lxrlocker/scale2models_wzy/wzy_e_55_track_r01.pth")
cudnn.benchmark = True best_wP = 0 start_epoch = 0 print('==> Preparing data..') BATCH_SIZE = 32 EPOCH = 30 INTERVAL = 200 model_name = 'efficient-all-frame-new' save_path = join("/root/data/wzy/checkpoint", model_name) train_dataloader, test_dataloader = get_dataloader(BATCH_SIZE) print('==> Building model..') net = EfficientNet.from_pretrained('efficientnet-b4', num_classes=2).cuda() net = nn.DataParallel(net).cuda() parser = argparse.ArgumentParser(description='PyTorch ImageNet Training') parser.add_argument('--lr-scheduler', default='cosine', type=str, metavar='SCHEDULER', help='LR scheduler (default: "step"') args = parser.parse_args() args.warmup_init_lr = 1e-5 args.lr = 2e-4 args.max_lr = args.lr args.warmup_updates = 5 * len(train_dataloader) args.max_update = EPOCH * len(train_dataloader)
import torch import cv2 from tqdm import tqdm import os import random from model.efficientnet.model import EfficientNet from torch.nn import functional as F random.seed(0) os.environ["CUDA_VISIBLE_DEVICES"] = "0" from lxrlocker.kagglefile.kaggle_lxr import MyModel, EfficientNet as lxrEfficientNet wzy_net = EfficientNet.from_name('efficientnet-b4', override_params={'num_classes': 2}).cuda().eval() wzy_net.load_state_dict(torch.load('/root/data/wzy/checkpoint/efficient-video/14.pth')) lxr_net = lxrEfficientNet.from_name('efficientnet-b4', override_params={'num_classes': 2}) lxr_net = MyModel(lxr_net).cuda().eval() lxr_net.load_state_dict(torch.load('/root/data/wzy/statistic/new_50partpai_tracking2_resize.pth')) from facenet_pytorch import MTCNN import numpy as np class FastMTCNN(object): def __init__(self, *args, **kwargs): self.mtcnn = MTCNN(*args, **kwargs) def __call__(self, frame): boxes, scores = self.mtcnn.detect([frame]) boxes, scores = boxes[0], scores[0] new_boxes = []