# torch.set_num_threads(12)
G = AEI_Net(c_id=512).to(device)
D = MultiscaleDiscriminator(input_nc=3,
n_layers=6,
norm_layer=torch.nn.InstanceNorm2d).to(device)
G.train()
D.train()
arcface = Backbone(50, 0.65, 'ir_se').to(device)
arcface.eval()
arcface.load_state_dict(torch.load('./id_model/model_ir_se50.pth',
map_location=device),
strict=False)
# weight_decay (float, optional)：权重衰减(如L2惩罚)(默认: 0)
opt_G = optim.Adam(G.parameters(),
lr=lr_G,
betas=(0, 0.999),
weight_decay=1e-5)
opt_D = optim.Adam(D.parameters(),
lr=lr_D,
betas=(0, 0.999),
weight_decay=1e-8)
G, opt_G = amp.initialize(G, opt_G, opt_level=optim_level)
D, opt_D = amp.initialize(D, opt_D, opt_level=optim_level)
try:
p_G = './saved_mask_landmarks_models/G_latest.pth'
p_D = './saved_mask_landmarks_models/D_latest.pth'
G.load_state_dict(torch.load(p_G, map_location=torch.device('cpu')),
# fine_tune_with_identity = False
device = torch.device('cuda')
# torch.set_num_threads(12)
G = AEI_Net(c_id=512).to(device)
D = MultiscaleDiscriminator(input_nc=3, n_layers=6, norm_layer=torch.nn.InstanceNorm2d).to(device)
G.train()
D.train()
arcface = Backbone(50, 0.6, 'ir_se').to(device)
arcface.eval()
arcface.load_state_dict(torch.load('./id_model/model_ir_se50.pth', map_location=device), strict=False)
opt_G = optim.Adam(G.parameters(), lr=lr_G, betas=(0, 0.999))
opt_D = optim.Adam(D.parameters(), lr=lr_D, betas=(0, 0.999))
G, opt_G = amp.initialize(G, opt_G, opt_level=optim_level)
D, opt_D = amp.initialize(D, opt_D, opt_level=optim_level)
try:
p_G = './saved_mask_models/G_latest.pth'
p_D = './saved_mask_models/D_latest.pth'
G.load_state_dict(torch.load(p_G, map_location=torch.device('cpu')), strict=False)
D.load_state_dict(torch.load(p_D, map_location=torch.device('cpu')), strict=False)
print('p_G : ',p_G)
print('p_D : ',p_D)
except Exception as e:
fine_tune_with_identity = False
device = torch.device('cuda')
# torch.set_num_threads(12)
G = AEI_Net(c_id=512).to(device)
D = MultiscaleDiscriminator(input_nc=3, n_layers=5, norm_layer=torch.nn.InstanceNorm2d).to(device)
G.train()
D.train()
arcface = Backbone(50, 0.6, 'ir_se').to(device)
arcface.eval()
arcface.load_state_dict(torch.load('./face_modules/model_ir_se50.pth', map_location=device), strict=False)
opt_G = optim.Adam(G.parameters(), lr=lr_G, betas=(0, 0.999), weight_decay=1e-4)
opt_D = optim.Adam(D.parameters(), lr=lr_D, betas=(0, 0.999), weight_decay=1e-4)
G, opt_G = amp.initialize(G, opt_G, opt_level=optim_level)
D, opt_D = amp.initialize(D, opt_D, opt_level=optim_level)
try:
G.load_state_dict(torch.load('./saved_models/G_latest.pth', map_location=torch.device('cpu')), strict=False)
D.load_state_dict(torch.load('./saved_models/D_latest.pth', map_location=torch.device('cpu')), strict=False)
except Exception as e:
print(e)
dataset = FaceEmbed([dataset_path], same_prob=0.5)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)