def train_style_transfer(args):
if not (args.train_data and args.valid_data):
print("must chose train_data and valid_data")
sys.exit()
# make dataset
trans = transforms.ToTensor()
train_dataset = FaceDataset(args.train_data, transform=trans)
label_dict = train_dataset.get_label_dict()
valid_dataset = FaceDataset(args.valid_data, transform=trans)
valid_dataset.give_label_dict(label_dict)
train_loader = data_utils.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=1)
valid_loader = data_utils.DataLoader(valid_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=1)
train_size = len(train_dataset)
valid_size = len(valid_dataset)
loaders = {"train": train_loader, "valid": valid_loader}
dataset_sizes = {"train": train_size, "valid": valid_size}
if args.gpu:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
device = torch.device("cpu")
# make network
if args.model_type == "VAE":
net = Autoencoder(train_dataset.label_num()).to(device)
optimizer = optim.Adam(net.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
best_model_wts = net.state_dict()
best_loss = 1e10
if args.generator_model and os.path.exists(args.generator_model):
net.load_state_dict(torch.load(args.generator_model))
elif args.model_type == "VAEGAN":
generator = Autoencoder(train_dataset.label_num()).to(device)
discriminator = Discriminator().to(device)
classifier = Classifier(train_dataset.label_num()).to(device)
generator_optimizer = optim.Adam(generator.parameters(),
lr=args.lr,
weight_decay=args.weight_decay)
discriminator_optimizer = optim.Adam(discriminator.parameters(),
lr=args.lr * 0.1,
weight_decay=args.weight_decay)
best_generator_wts = generator.state_dict()
best_discriminator_wts = discriminator.state_dict()
best_generator_loss = 1e10
best_discriminator_loss = 1e10
if args.generator_model and os.path.exists(args.generator_model):
generator.load_state_dict(torch.load(args.generator_model))
if args.discriminator_model and os.path.exists(
args.discriminator_model):
discriminator.load_state_dict(torch.load(args.discriminator_model))
if args.classifier_model:
classifier.load_state_dict(torch.load(args.classifier_model))
# make loss function and optimizer
criterion = nn.BCELoss(reduction="sum")
classifier_criterion = nn.CrossEntropyLoss(reduction="sum")
# initialize loss
loss_history = {"train": [], "valid": []}
# start training
start_time = time.time()
for epoch in range(args.epochs):
print("epoch {}".format(epoch + 1))
for phase in ["train", "valid"]:
if phase == "train":
if args.model_type == "VAE":
net.train(True)
elif args.model_type == "VAEGAN":
generator.train(True)
discriminator.train(True)
else:
if args.model_type == "VAE":
net.train(False)
elif args.model_type == "VAEGAN":
generator.train(False)
discriminator.train(False)
# initialize running loss
generator_running_loss = 0.0
discriminator_running_loss = 0.0
for i, data in enumerate(loaders[phase]):
inputs, label = data
# wrap the in valiables
if phase == "train":
inputs = Variable(inputs).to(device)
label = Variable(label).to(device)
torch.set_grad_enabled(True)
else:
inputs = Variable(inputs).to(device)
label = Variable(label).to(device)
torch.set_grad_enabled(False)
# zero gradients
if args.model_type == "VAE":
optimizer.zero_grad()
mu, var, outputs = net(inputs, label)
loss = loss_func(inputs, outputs, mu, var)
if phase == "train":
loss.backward()
optimizer.step()
generator_running_loss += loss.item()
elif args.model_type == "VAEGAN":
real_label = Variable(
torch.ones((inputs.size()[0], 1), dtype=torch.float) -
0.2 * (torch.rand(inputs.size()[0], 1))).to(device)
fake_label = Variable(
torch.zeros((inputs.size()[0], 1), dtype=torch.float) +
0.2 * (torch.rand(inputs.size()[0], 1))).to(device)
discriminator_optimizer.zero_grad()
real_pred = discriminator(inputs)
real_loss = criterion(real_pred, real_label)
random_index = np.random.randint(0,
train_dataset.label_num(),
inputs.size()[0])
generate_label = Variable(
torch.zeros_like(label)).to(device)
for i, index in enumerate(random_index):
generate_label[i][index] = 1
mu, var, outputs = generator(inputs, label)
fake_pred = discriminator(outputs.detach())
fake_loss = criterion(fake_pred, fake_label)
discriminator_loss = real_loss + fake_loss
if phase == "train":
discriminator_loss.backward()
discriminator_optimizer.step()
generator_optimizer.zero_grad()
#class_loss = classifier_criterion(classifier(outputs), torch.max(label, 1)[1])
dis_loss = criterion(discriminator(outputs), real_label)
gen_loss = loss_func(inputs, outputs, mu, var)
generator_loss = dis_loss + gen_loss
if phase == "train":
generator_loss.backward()
generator_optimizer.step()
discriminator_running_loss += discriminator_loss.item()
generator_running_loss += generator_loss.item()
if args.model_type == "VAE":
epoch_loss = generator_running_loss / dataset_sizes[
phase] * args.batch_size
loss_history[phase].append(epoch_loss)
print("{} loss {:.4f}".format(phase, epoch_loss))
if phase == "valid" and epoch_loss < best_loss:
best_model_wts = net.state_dict()
best_loss = epoch_loss
elif args.model_type == "VAEGAN":
epoch_generator_loss = generator_running_loss / dataset_sizes[
phase] * args.batch_size
epoch_discriminator_loss = discriminator_running_loss / dataset_sizes[
phase] * args.batch_size
print("{} generator loss {:.4f}".format(
phase, epoch_generator_loss))
print("{} discriminator loss {:.4f}".format(
phase, epoch_discriminator_loss))
if phase == "valid" and epoch_generator_loss < best_generator_loss:
best_generator_wts = generator.state_dict()
best_generator_loss = epoch_generator_loss
if phase == "valid" and epoch_discriminator_loss < best_discriminator_loss:
best_discriminator_wts = discriminator.state_dict()
best_generator_loss = epoch_discriminator_loss
elapsed_time = time.time() - start_time
print("training complete in {:.0f}s".format(elapsed_time))
if args.model_type == "VAE":
net.load_state_dict(best_model_wts)
return net, label_dict
elif args.model_type == "VAEGAN":
generator.load_state_dict(best_generator_wts)
discriminator.load_state_dict(best_discriminator_wts)
return (generator, discriminator), label_dict
N_val_btot = sdat_val.get_Nbatches_tot()
print('Training Batches: ', N_train_btot)
print('Validation Batches: ', N_val_btot)
# ----------- START TRAINING LOOP ---------------------
for i_epoch in range(num_epochs):
d_train_wloss = 0.0
d_train_gploss = 0.0
g_train_loss = 0.0
d_val_wloss = 0.0
d_val_gploss = 0.0
g_val_loss = 0.0
G.train()
D.train()
for i_batch in range(N_train_btot):
for i_c in range(n_critic):
# ---------- DISCRIMINATOR STEP ---------------
(data_b, i_vc) = sdat_train.get_rand_batch()
real_wfs = torch.from_numpy(data_b).float()
i_vc = [torch.from_numpy(i_v).float() for i_v in i_vc]
Nsamp = real_wfs.size(0)
if cuda:
real_wfs = real_wfs.cuda()
i_vc = [i_v.cuda() for i_v in i_vc]
d_optimizer.zero_grad()
z = noise(Nbatch, z_size)
z = torch.from_numpy(z).float()
if cuda:
fixed_noise = torch.randn(128, latent_size, 1, 1, device=device)
history = {
'lossD': [],
'lossG': [],
'valLossD': [],
'avgStart': [],
'avgEnd': [],
'valid': [],
'validHist': [],
'imgs': [],
'epoch': []
}
for epoch in range(n_epochs):
netD.train()
netG.train()
for i, data in enumerate(loader['train']):
x_real = data['moves'].type(Tensor)
valid = Tensor(x_real.size(0), 1).fill_(1)
fake = Tensor(x_real.size(0), 1).fill_(0)
if epoch == 0 and i == 0:
viz.images(x_real[:16].flip(2),
opts=dict(title='Epoch' + str(epoch),
width=1000,
height=250))
netD.zero_grad()