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 = []