shutil.rmtree(tensorboard_path)
os.makedirs(tensorboard_path)
else:
os.makedirs(tensorboard_path)
torch.save(
{
'epoch': epoch,
'lossesG': lossesG,
'lossesD': lossesD,
'E_state_dict': E.state_dict(),
'G_state_dict': G.state_dict(),
# 'D_state_dict': D.state_dict(),
'optimizerG_state_dict': optimizerG.state_dict(),
# 'optimizerD_state_dict': optimizerD.state_dict(),
'num_vid': dataset.__len__(),
'i_batch': i_batch
},
path_to_chkpt)
print('...Done')
"""Loading from past checkpoint"""
checkpoint = torch.load(path_to_chkpt, map_location=cpu)
E.load_state_dict(checkpoint['E_state_dict'])
G.load_state_dict(checkpoint['G_state_dict'])
# D.load_state_dict(checkpoint['D_state_dict'])
optimizerG.load_state_dict(checkpoint['optimizerG_state_dict'])
# optimizerD.load_state_dict(checkpoint['optimizerD_state_dict'])
epochCurrent = checkpoint['epoch']
lossesG = checkpoint['lossesG']
lossesD = checkpoint['lossesD']
num_vid = checkpoint['num_vid']
device = torch.device("cuda:0")
cpu = torch.device("cpu")
tensorboard_path = './experiment/tensorboard'
path_to_chkpt = './experiment/model_weights_self_train.tar'
path_to_backup = './experiment/backup_model_weights.tar'
path_to_mp4 = "/home/cxu-serve/p1/common/voxceleb/test/video/sample"
VGGFace_body_path = '/home/cxu-serve/p1/common/vggface/new/Pytorch_VGGFACE_IR.py'
VGGFace_weight_path = '/home/cxu-serve/p1/common/vggface/new/Pytorch_VGGFACE.pth'
dataset = VidDataSet(K=8, path_to_mp4=path_to_mp4, device=device)
dataLoader = DataLoader(dataset, batch_size=2, shuffle=True)
G = Generator(224).to(device)
E = Embedder(224).to(device)
D = Discriminator(dataset.__len__()).to(device)
G.train()
E.train()
D.train()
optimizerG = optim.Adam(params=list(E.parameters()) + list(G.parameters()),
lr=5e-5)
optimizerD = optim.Adam(params=D.parameters(), lr=2e-4)
"""Criterion"""
criterionG = LossG(VGGFace_body_path=VGGFace_body_path,
VGGFace_weight_path=VGGFace_weight_path,
device=device)
criterionDreal = LossDSCreal()
criterionDfake = LossDSCfake()
"""Training init"""
os.environ['CUDA_VISIBLE_DEVICES'] = '1,0'
device = torch.device("cuda")
cpu = torch.device("cpu")
path_to_chkpt = 'model_weights.tar'
path_to_backup = 'backup_model_weights.tar'
dataset = VidDataSet(
K=8,
path_to_mp4='/data2/hao66/dataset/voxceleb1/unzippedFaces',
device=device)
print('# of videos: ', len(dataset))
dataLoader = DataLoader(dataset, batch_size=2, shuffle=True)
G = torch.nn.DataParallel(Generator(224)).to(device)
E = torch.nn.DataParallel(Embedder(224)).to(device)
D = torch.nn.DataParallel(Discriminator(dataset.__len__())).to(device)
G.train()
E.train()
D.train()
optimizerG = optim.Adam(params=list(E.parameters()) + list(G.parameters()),
lr=5e-5)
optimizerD = optim.Adam(params=D.parameters(), lr=2e-4)
"""Criterion"""
criterionG = LossG(device=device)
criterionDreal = LossDSCreal()
criterionDfake = LossDSCfake()
"""Training init"""
epochCurrent = epoch = i_batch = 0
lossesG = []