head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
if cfg.TRAIN.num_gpus > 1 and cfg.TRAIN.use_gpu:
net = torch.nn.DataParallel(net).cuda()
else:
net = net.cuda()
cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(),
lr=cfg.TRAIN.LR.initial_lr,
momentum=cfg.TRAIN.LR.momentum,
weight_decay=cfg.TRAIN.LR.weight_decay)
criterion = MultiBoxLoss(cfg)
anchor_generator = AnchorGenerator(cfg)
with torch.no_grad():
anchors = anchor_generator.generate_anchors()
anchors = anchors.cuda()
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
#--------------------------------------------#
# 该部分代码只用于看网络结构，并非测试代码
#--------------------------------------------#
import torch
from torchsummary import summary
from nets.retinaface import RetinaFace
from utils.config import cfg_mnet
if __name__ == '__main__':
# 需要使用device来指定网络在GPU还是CPU运行
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = RetinaFace(cfg_mnet).to(device)
print('# generator parameters:',
sum(param.numel() for param in model.parameters()))
summary(model, input_size=(3, 2150, 2150))