def multiple_demo(load_fnames, save_fname="multi", reg_a=10., reg_d=10.):
# assignment: section 6
bfm = h5py.File(MODELS_PATH + BFM_FNAME, "r")
bfm_params, color, triangles = utils.read_bfm(bfm)
lms = utils.read_landmarks(MODELS_PATH + LM_FNAME) # landmark annotations
N = len(load_fnames) # number of images to be loaded
imgs_real = [utils.read_image(IMAGE_PATH + fname) for fname in load_fnames] # load all images
hs = [np.size(img, 0) for img in imgs_real] # store all heights
ws = [np.size(img, 1) for img in imgs_real] # store all widths
lms_real = [torch.from_numpy(detect_landmark(img)) for img in imgs_real] # detect all ground truth landmarks
lms_real_flip = [utils.flip_ycoords(lms_real[i], H=hs[i]) for i in range(N)] # flip y axis
alpha, deltas, rotations, translations, loss = multiple.estimate_params((bfm_params, color, triangles),
lms, lms_real_flip, hs=hs, ws=ws,
reg_a=reg_a, reg_d=reg_d)
utils.save_loss(loss, save_fname=save_fname + "_loss.pdf")
# save results for each image
for i in range(N):
print(load_fnames[i] + ":") # print stats for each image (alpha is the same for each img)
utils.print_stats(alpha, deltas[i], rotations[i], translations[i])
G = morph.compute_G(bfm_params, alpha=alpha, delta=deltas[i])
G_transformed = pinhole.transform(G, rotations[i], translations[i])
G_pinhole = pinhole.camera_model(G, rotations[i], translations[i], h=hs[i], w=ws[i])
color = texture.get_color(imgs_real[i], G_pinhole[:, :2])
print("Rendering...")
img_pred = utils.get_image(G_pinhole, color, triangles, h=hs[i], w=ws[i])
utils.show_face(img_pred)
utils.flip_y()
plt.savefig(PINHOLE_PATH + save_fname + str(i) + ".pdf")
plt.close()
save_obj(OBJ_3D_PATH + save_fname + str(i) + "_3d.obj", G_transformed, color, triangles)
save_obj(OBJ_2D_PATH + save_fname + str(i) + "_2d.obj", G_pinhole, color, triangles)
lm_pred_flip = utils.get_landmarks(G_pinhole[:, :2], lms)
lm_pred = utils.flip_ycoords(lm_pred_flip, H=hs[i])
utils.show_face(imgs_real[i], white_background=False)
utils.show_landmarks(lms_real[i], indices=False, label="ground-truth")
try:
utils.show_landmarks(lm_pred, indices=False, label="model")
except TypeError:
print("... unable to show predicted landmarks")
plt.savefig(PINHOLE_PATH + save_fname + str(i) + "_lm.pdf")
plt.close()
def trainer(args):
"""Training loop.
Handles model, optimizer, loss, and sampler generation.
Handles data loading. Handles i/o and checkpoint loading.
Parameters
----------
conf : dict
Miscellaneous parameters
"""
############### Dataset ########################
loader = data.load_denoising(args.data_root,
train=True,
batch_size=args.batch_size,
transform=None)
val_loader = data.load_denoising(args.data_root,
train=False,
batch_size=args.batch_size,
transform=None)
##################################################
##### Model, Optimizer, Loss ############
num_classes = 30
network = model.Unet(args.unet_hidden, num_classes).to(args.device)
optimizer = torch.optim.Adam(network.parameters(), lr=args.lr)
criterion = nn.CrossEntropyLoss()
##################################################
############## Training loop #####################
for epoch in range(args.load_checkpoint + 1, args.epochs + 1):
print('\nEpoch %d/%d' % (epoch, args.epochs))
# Train
network, optimizer, train_epoch_loss, train_epoch_dice = train(
network, loader, criterion, optimizer, args.device)
# Validate
network, val_epoch_loss, val_epoch_dice = validate(
network, val_loader, criterion, args.device)
# Save checkpoints
utils.save_checkpoint(epoch, network.state_dict(), optimizer.state_dict(), \
train_epoch_loss, val_epoch_loss, args.filename, args.log_interval)
utils.save_loss(epoch, train_epoch_loss, val_epoch_loss, args.run_dir)
utils.save_loss(epoch, train_epoch_dice, val_epoch_dice, args.run_dir,
'dice')
print("Epoch {}: test loss: {:.6f}, test dice: {:.6f}".format(
epoch, val_epoch_loss, val_epoch_dice))
def texture_demo(load_fname="yke_neutral.jpeg", save_fname="texture", reg_a=10., reg_d=10., idx=0):
# assignment: section 5
bfm = h5py.File(MODELS_PATH + BFM_FNAME, "r")
bfm_params, color, triangles = utils.read_bfm(bfm)
lms = utils.read_landmarks(MODELS_PATH + LM_FNAME) # landmark annotations
img_real = utils.read_image(IMAGE_PATH + load_fname) # load image of face we want to reconstruct
h, w, _ = np.shape(img_real)
lm_real = torch.from_numpy(detect_landmark(img_real)) # detect ground-truth landmarks
lm_real_flip = utils.flip_ycoords(lm_real, H=h) # flip y axis because img is upside down compared to pinhole output
alpha, delta, rotation, translation, loss = latent.estimate_params((bfm_params, color, triangles),
lms, lm_real_flip, h=h, w=w,
reg_a=reg_a, reg_d=reg_d)
utils.print_stats(alpha, delta, rotation, translation) # latent params statistics
utils.save_loss(loss, save_fname=save_fname + str(idx) + "_loss.pdf")
G = morph.compute_G(bfm_params, alpha=alpha, delta=delta)
G_pinhole = pinhole.camera_model(G, rotation, translation, h=h, w=w)
color = texture.get_color(img_real, G_pinhole[:, :2]) # obtain vertex colors from provided image
save_obj(OBJ_3D_PATH + save_fname + str(idx) + "_3d.obj", G, color, triangles)
save_obj(OBJ_2D_PATH + save_fname + str(idx) + "_2d.obj", G_pinhole, color, triangles)
print("Rendering...")
img_pred = utils.get_image(G_pinhole, color, triangles, h=h, w=w)
utils.show_face(img_pred)
utils.flip_y()
plt.savefig(PINHOLE_PATH + save_fname + str(idx) + ".pdf")
plt.close()
lm_pred_flip = utils.get_landmarks(G_pinhole[:, :2], lms)
lm_pred = utils.flip_ycoords(lm_pred_flip, H=h)
utils.show_face(img_real, white_background=False)
utils.show_landmarks(lm_real, indices=False, label="ground-truth")
try:
utils.show_landmarks(lm_pred, indices=False, label="model")
except TypeError:
print("... unable to show predicted landmarks")
plt.savefig(PINHOLE_PATH + save_fname + str(idx) + "_lm.pdf")
plt.close()
args.save_dir, "iris", "layers{}_eps{}_eta{}"
"".format(args.layers, args.eps, args.eta))
os.makedirs(model_dir, exist_ok=True)
utils.save_params(model_dir, vars(args))
# data setup
X_train, y_train, X_test, y_test, num_classes = dataset.load_Iris(0.3)
# model setup
layers = [Vector(args.layers, eta=args.eta, eps=args.eps)]
model = Architecture(layers, model_dir, num_classes)
# train/test pass
print("Forward pass - train features")
Z_train = model(X_train, y_train)
utils.save_loss(model.loss_dict, model_dir, "train")
print("Forward pass - test features")
Z_test = model(X_test)
utils.save_loss(model.loss_dict, model_dir, "test")
# save features
utils.save_features(model_dir, "X_train", X_train, y_train)
utils.save_features(model_dir, "X_test", X_test, y_test)
utils.save_features(model_dir, "Z_train", Z_train, y_train)
utils.save_features(model_dir, "Z_test", Z_test, y_test)
# evaluation train
_, acc_svm = evaluate.svm(Z_train, y_train, Z_train, y_train)
acc_knn = evaluate.knn(Z_train, y_train, Z_train, y_train, k=5)
acc_svd = evaluate.nearsub(Z_train, y_train, Z_train, y_train, n_comp=1)
acc = {"svm": acc_svm, "knn": acc_knn, "nearsub-svd": acc_svd}
def main():
clear_loss()
parser = get_parser()
parsed = parser.parse_args()
assert ((parsed.output_dir is None and parsed.save_step is None) or
(parsed.output_dir is not None and parsed.save_step is not None)), "Save step and output directory must be " \
"null at the same time or not null at the same time"
ds_type = parsed.dataset
if ds_type == 'cifar10':
dataset = CIFAR10Dataset()
dataset.process()
img_shape = [32, 32, 3]
elif ds_type == 'mnist':
dataset = MNISTDataset()
dataset.process()
img_shape = [28, 28, 1]
elif parsed.dataset == 'fashion':
dataset = FashionDataset()
dataset.process()
img_shape = [28, 28, 1]
elif parsed.dataset == 'stl10':
dataset = STLDataset(is_ae=True)
dataset.process()
img_shape = [96, 96, 3]
else:
print("Unknown dataset")
exit()
layers = parse_layers(parsed.layer_str)
fc_size = parsed.fc_layers
sess = tf.Session()
swwae = SWWAE(sess,
img_shape,
'autoencode',
layers,
learning_rate=parsed.learning_rate,
lambda_rec=parsed.lambda_rec,
lambda_M=parsed.lambda_M,
dtype=tf.float32,
tensorboard_id=parsed.tensorboard_id,
encoder_train=True,
rep_size=fc_size,
batch_size=parsed.batch_size,
sparsity=parsed.sparsity,
beta=parsed.beta)
if parsed.rest_dir is not None:
swwae.restore(parsed.rest_dir)
X_test, _ = dataset.get_batches(parsed.batch_size, train=False)
test_steps = len(X_test)
print("Preprocessing")
datagen = ImageDataGenerator(
featurewise_center=False, # set input mean to 0 over the dataset
samplewise_center=False, # set each sample mean to 0
featurewise_std_normalization=
False, # divide inputs by std of the dataset
samplewise_std_normalization=False, # divide each input by its std
zca_whitening=False, # apply ZCA whitening
rotation_range=
0.0, # randomly rotate images in the range (degrees, 0 to 180)
width_shift_range=
0.1, # randomly shift images horizontally (fraction of total width)
height_shift_range=
0.1, # randomly shift images vertically (fraction of total height)
horizontal_flip=True, # randomly flip images
vertical_flip=False) # randomly flip images
datagen.fit(dataset.training_data)
train_steps = int(len(dataset.training_data) / parsed.batch_size)
print("Started training.\nTrain steps: {}".format(train_steps))
for e in range(parsed.num_epochs):
total_loss = 0.0
epoch_loss = 0.0
batches = 0
for x_batch in datagen.flow(dataset.training_data,
batch_size=parsed.batch_size):
loss, global_step = swwae.train(x_batch)
batches += 1
total_loss += loss
epoch_loss += loss
if (batches + 1) % parsed.info_step == 0:
avg_loss = total_loss / parsed.info_step
save_loss(avg_loss)
for test_step in range(test_steps):
X_test_step = X_test[test_step]
swwae.eval(input=X_test_step)
#print("Train epoch {}:\n\tstep {}\n\tavg. L2 Loss: {}".format(e + 1, step + 1, avg_loss),
# flush=True)
total_loss = 0.0
if parsed.save_step is not None:
if (global_step + 1) % parsed.save_step == 0:
swwae.save(path=parsed.output_dir)
if batches >= train_steps:
break
print("Train epoch {}: avg. loss: {}".format(e + 1,
epoch_loss / train_steps),
flush=True)
if parsed.output_dir is not None:
swwae.save(path=parsed.output_dir)
print("Starting test..")
total_loss = 0.0
for test_step in range(test_steps):
X_test_step = X_test[test_step]
loss = swwae.eval(input=X_test_step)
total_loss += loss
print("Test average loss: {}".format(total_loss / test_steps))
X_test, y_test = X_test.to(device), y_test.to(device)
## Architecture
net = L.load_architecture(params['data'], params['arch'])
net = utils.load_ckpt(args.model_dir, 'model', net)
net = net.to(device)
## Forward
with torch.no_grad():
print('train')
Z_train = net.batch_forward(X_train,
batch_size=args.batch_size,
loss=args.loss,
device=device)
X_train, y_train, Z_train = F.to_cpu(X_train, y_train, Z_train)
utils.save_loss(eval_dir, f'train', net.get_loss())
print('test')
Z_test = net.batch_forward(X_test,
batch_size=args.batch_size,
loss=args.loss,
device=device)
X_test, y_test, Z_test = F.to_cpu(X_test, y_test, Z_test)
utils.save_loss(eval_dir, f'test', net.get_loss())
## Normalize
X_train = F.normalize(X_train.flatten(1))
X_test = F.normalize(X_test.flatten(1))
Z_train = F.normalize(Z_train.flatten(1))
Z_test = F.normalize(Z_test.flatten(1))
#print "upscore2 params", solver.net.params['upscore2_1'][0].data[0,0,:,:]
# take snapshot
if solver.iter % par.train_snapshot_iters == 0:
filename = par.train_save_path + "corrNetwork_" + str(
solver.iter) + ".caffemodel"
solver.net.save(filename)
# store contrastive loss
loss_contr[solver.iter] = solver.net.blobs['contr_loss'].data
# visualize the contrastive loss and save the pair distances
if solver.iter % par.verbose_iters == 0:
vis.plot_loss(loss_contr, solver.iter, 1, par.train_save_path,
"contr")
utils.save_pairs_dist(par.train_save_path, solver.net,
solver.iter)
# save all loss values in a txt file
utils.save_loss(loss_contr[1:], par.train_save_path, "contr")
elif par.phase == "TEST_SCENE":
corrNet.generateDeployPrototxtFile() # generate the deploy prototxt
caffemodel = par.test_model_path + "corrNetwork_" + str(
par.test_iter) + ".caffemodel"
print "Using caffemodel: " + caffemodel
flann = FLANN()
#scene_id = 2 # 4, 6
#nfiles = 718 # 704 # 441
scene_id = [2, 4, 6]
nfiles = [718, 704, 441]
for sc in range(len(scene_id)):
putative_list, precision_list, recall_list = [], [], []
for i in range(0, nfiles[sc], 5):
img1, fr_name1, _ = gmu_input_data.pick_im(scene_id[sc], i)
tf.concat(greedy_seqs.values, axis=0).numpy())
train_inds = list(
tf.concat(batch[-1].values, axis=0).numpy().squeeze())
articles = [article[x] for x in train_inds]
gt_summaries = [summary[x] for x in train_inds]
examples_oovs = [oovs[x] for x in train_inds]
scores, summaries, time_step_masks = env.get_rewards(
gt_summaries, train_sums, examples_oovs)
save_examples(examples_folder, articles, gt_summaries, summaries,
epoch, batch_n, 'train')
save_scores(metrics_folder, scores, 'train')
mean_epoch_loss = np.mean(losses)
losses = []
save_loss(metrics_folder, mean_epoch_loss, 'train')
val_losses = []
val_sums = []
val_inds = []
val_iterator = iter(val_dist_dataset)
for val_batch_n in range(1, min(10, batches_per_epoch)):
batch = next(val_iterator)
loss, greedy_seqs = distributed_step(batch, 'val')
val_losses.append(loss)
with tf.device('CPU'):
val_sums += list(
tf.concat(greedy_seqs.values, axis=0).numpy())
val_inds += list(
tf.concat(batch[-1].values, axis=0).numpy().squeeze())
def train(self):
r_total_loss = []
d_total_loss = []
g_total_loss = []
for epoch in range(args.n_epochs):
r_epoch_loss = 0
d_epoch_loss = 0
g_epoch_loss = 0
for batch, (imgs, _) in enumerate(self.images_loader):
with torch.no_grad():
imgs = Variable(imgs.type(Tensor))
# Train Autoencoder - Reconstruction
self.optimizer_R.zero_grad()
if args.architecture == 'ResNet':
encoded_imgs = self.encoder(imgs)
decoded_imgs = self.decoder(encoded_imgs)
r_loss = self.reconstruct_loss(decoded_imgs, imgs)
r_loss.backward(retain_graph=True)
self.optimizer_R.step()
# Train Discriminator
for k in range(args.k):
self.optimizer_D.zero_grad()
with torch.no_grad():
# sample noise
z = Variable(
Tensor(
np.random.normal(
0, args.s_sd,
(imgs.shape[0], args.latent_dim))))
# measure discriminator's ability to classify real from generated samples
d_real = self.discriminator(utils.gaussian(
z, 0, args.n_sd))
d_fake = self.discriminator(
utils.gaussian(encoded_imgs.detach(), 0, args.n_sd))
d_loss = -(torch.mean(d_real) - torch.mean(d_fake))
d_loss.backward()
self.optimizer_D.step()
# clip discriminator's weights
for p in self.discriminator.parameters():
p.data.clamp_(-args.clip_value, args.clip_value)
# Train Generator
self.optimizer_G.zero_grad()
d_fake = self.discriminator(
utils.gaussian(encoded_imgs, 0, args.n_sd))
g_loss = -torch.mean(d_fake)
g_loss.backward()
self.optimizer_G.step()
# write losses to files
with torch.no_grad():
utils.save_loss(epoch + 1, batch + 1,
len(self.images_loader), r_loss.item(),
'reconstruction')
utils.save_loss(epoch + 1, batch + 1,
len(self.images_loader), d_loss.item(),
'discriminator')
utils.save_loss(epoch + 1, batch + 1,
len(self.images_loader), g_loss.item(),
'generator')
r_epoch_loss += r_loss.item()
d_epoch_loss += d_loss.item()
g_epoch_loss += g_loss.item()
batches_done = epoch * len(self.images_loader) + batch
if batches_done % args.sample_interval == 0:
utils.save_images(imgs,
'real',
n_row=int(math.sqrt(
args.batch_size)),
batches_done=batches_done)
utils.save_images(decoded_imgs.detach(),
'recon',
n_row=int(math.sqrt(
args.batch_size)),
batches_done=batches_done)
print('Images and their reconstructions saved')
with torch.no_grad():
r_total_loss.append(r_epoch_loss / len(self.images_loader))
d_total_loss.append(d_epoch_loss / len(self.images_loader))
g_total_loss.append(g_epoch_loss / len(self.images_loader))
# save loss plots
utils.save_plot(epoch, r_total_loss, 'Reconstruction')
utils.save_plot(epoch, d_total_loss, 'Discriminator')
utils.save_plot(epoch, g_total_loss, 'Generator')
print('Plots saved')
# save images generated from fixed noise
gen_random = self.decoder(self.fixed_noise)
utils.save_images(gen_random.detach(),
'generated_fixed',
n_row=10,
batches_done=batches_done)
print('Images generated from fixed noise saved')
if epoch % 5 == 0:
# save images generated from random noise
z = Variable(
Tensor(
np.random.normal(0, args.s_sd,
(100, args.latent_dim))))
gen_random = self.decoder(z)
utils.save_images(gen_random.detach(),
'generated_random',
n_row=10,
batches_done=batches_done)
print('Images generated from random noise saved')
#utils.save_manifold(batches_done, self.decoder)
#print('Manifold saved')
# save models
torch.save(self.encoder.state_dict(),
'{}/encoder.pth'.format(args.folder_name))
torch.save(self.decoder.state_dict(),
'{}/decoder.pth'.format(args.folder_name))
torch.save(self.discriminator.state_dict(),
'{}/discriminator.pth'.format(args.folder_name))
print('Models saved')
# decay learning rate
self.scheduler_R.step()
print('Epoch:', epoch, 'R LR:', self.scheduler_R.get_lr())
self.scheduler_D.step()
print('Epoch:', epoch, 'D LR:', self.scheduler_D.get_lr())
self.scheduler_G.step()
print('Epoch:', epoch, 'G LR:', self.scheduler_G.get_lr())
def train(net, epochs, batch_size, print_every=50):
# mnist datset loader
mnist_dir = 'mnist'
mnist = input_data.read_data_sets(mnist_dir, one_hot=True)
steps = 0
margin_losses = []
recontruction_losses = []
total_losses = []
accuracies = []
start_time = time.time()
with tf.Session() as sess:
# reset tensorflow variables
sess.run(tf.global_variables_initializer())
# start training
for e in range(epochs):
for ii in range(mnist.train.num_examples // batch_size):
# get training data
batch_x, batch_y = mnist.train.next_batch(batch_size)
# reshape input
batch_x = np.reshape(batch_x, (-1, 28, 28, 1))
fd = {net.inputs_x: batch_x, net.inputs_y: batch_y}
# Run optimizers
_ = sess.run(net.train_opt, feed_dict=fd)
# evaluate losses
if steps % print_every == 0:
margin_loss = net.margin_loss.eval(fd)
recon_loss = net.recon_loss.eval(fd)
total_loss = net.total_loss.eval(fd)
# compute current accuracy
accuracy = compute_accuracy(sess, net, mnist.validation)
print(
"Epoch {}/{}...".format(e + 1, epochs),
"Margin Loss: {:.4f}...".format(margin_loss),
"Reconstruction Loss: {:.4f}...".format(recon_loss),
"Total Loss: {:.4f}...".format(total_loss),
"Epoch {}/{}... Accuracy: {:.05f}".format(
e + 1, epochs, accuracy))
# save losses & accuracies
margin_losses.append(margin_loss)
recontruction_losses.append(recon_loss)
total_losses.append(total_loss)
accuracies.append(accuracy)
steps += 1
# get reconstructed results
recon_result_fn = 'recon_{:03d}.png'.format(e + 1)
save_reconstruction_results(net, mnist.validation, recon_result_fn)
# test final accuracy and reconstructions
final_test_accuracy = compute_accuracy(sess, net, mnist.test)
print('Final test accuracy: {:.4f}'.format(final_test_accuracy))
final_recon_fn = 'final_recon.png'
save_reconstruction_results(net, mnist.test, final_recon_fn)
end_time = time.time()
elapsed_time = end_time - start_time
print('Elapsed: {:05.03f} s'.format(elapsed_time))
# save losses as image
margin_losses_fn = 'margin-loss.png'
recon_losses_fn = 'reconstruction-loss.png'
total_losses_fn = 'total-loss.png'
accuracy_fn = 'accuracy.png'
utils.save_loss(margin_losses, 'Margin-loss', margin_losses_fn)
utils.save_loss(recontruction_losses, 'Reconstruction-loss',
recon_losses_fn)
utils.save_loss(total_losses, 'Total-loss', total_losses_fn)
utils.save_loss(accuracies, 'Accuracy', accuracy_fn)
return
def train_and_evaluate(model: nn.Module,
discriminator:nn.Module,
train_loader: DataLoader,
valid_loader: DataLoader,
test_loader: DataLoader,
optimizer_G,
optimizer_D,
adversarial_loss,
params: utils.Params,
restore_file: str = None) -> None:
'''Train the model and evaluate every epoch.
Args:
model: (torch.nn.Module) the Deep AR model
train_loader: load train data and labels
test_loader: load test data and labels
optimizer: (torch.optim) optimizer for parameters of model
loss_fn: a function that takes outputs and labels per timestep, and then computes the loss for the batch
params: (Params) hyperparameters
restore_file: (string) optional- name of file to restore from (without its extension .pth.tar)
'''
early_stopping = EarlyStopping(patience=100, verbose=True)
# reload weights from restore_file if specified
if restore_file is not None:
restore_path = os.path.join(params.model_dir, restore_file + '.pth.tar')
logger.info('Restoring parameters from {}'.format(restore_path))
utils.load_checkpoint(restore_path, model, optimizer_G)
logger.info('begin training and evaluation')
best_valid_q50 = float('inf')
best_valid_q90 = float('inf')
best_test_q50 = float('inf')
best_test_q90 = float('inf')
best_MAPE = float('inf')
train_len = len(train_loader)
test_len = len(test_loader)
q50_summary = np.zeros(params.num_epochs)
q90_summary = np.zeros(params.num_epochs)
MAPE_summary = np.zeros(params.num_epochs)
q50_valid = np.zeros(params.num_epochs)
q90_valid = np.zeros(params.num_epochs)
MAPE_valid = np.zeros(params.num_epochs)
loss_summary = np.zeros((train_len * params.num_epochs))
loss_test = np.zeros((test_len * params.num_epochs))
d_loss_summary = np.zeros((train_len * params.num_epochs))
valid_loss = []
logger.info("My Transformer have {} paramerters in total".format(sum(x.numel() for x in model.parameters())))
for epoch in range(params.num_epochs):
logger.info('Epoch {}/{}'.format(epoch + 1, params.num_epochs))
loss_summary[epoch * train_len:(epoch + 1) * train_len], d_loss_summary[epoch * train_len:(epoch + 1) * train_len] = train(model, discriminator,optimizer_G, optimizer_D, adversarial_loss, train_loader,
valid_loader, params, epoch)
test_metrics = evaluate(model, test_loader, params, epoch)
valid_metrics = evaluate(model, valid_loader, params, epoch)
loss_test[epoch * test_len:(epoch + 1) * test_len] = test_metrics['loss'].cpu()
q50_valid[epoch] = valid_metrics['q50']
q90_valid[epoch] = valid_metrics['q90']
MAPE_valid[epoch] = valid_metrics['MAPE']
q50_summary[epoch] = test_metrics['q50']
q90_summary[epoch] = test_metrics['q90']
MAPE_summary[epoch] = test_metrics['MAPE']
valid_loss.append(valid_metrics['q50'])
#is_best = q90_summary[epoch] <= best_test_q90
is_best = q50_valid[epoch] <= best_valid_q50
# Save weights
utils.save_checkpoint({'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optim_dict': optimizer_G.state_dict()},
epoch=epoch,
is_best=is_best,
checkpoint=params.model_dir)
if is_best:
logger.info('- Found new best Q90/Q50')
best_valid_q50 = q50_summary[epoch]
best_valid_q50 = q50_valid[epoch]
best_json_path = os.path.join(params.model_dir, 'metrics_test_best_weights.json')
utils.save_dict_to_json(test_metrics, best_json_path)
utils.save_loss(loss_summary[epoch * train_len:(epoch + 1) * train_len], args.dataset + '_' + str(epoch) +'-th_epoch_loss', params.plot_dir)
utils.save_loss(loss_test[epoch * test_len:(epoch + 1) * test_len], args.dataset + '_' + str(epoch) +'-th_epoch_test_loss', params.plot_dir)
last_json_path = os.path.join(params.model_dir, 'metrics_test_last_weights.json')
utils.save_dict_to_json(test_metrics, last_json_path)
early_stopping(valid_loss[-1], model)
if early_stopping.early_stop:
print("Early stopping")
# save weights
utils.save_checkpoint({'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optim_dict': optimizer.state_dict()},
filepath=params.model_dir)
break
if args.save_best:
f = open('./param_search.txt', 'w')
f.write('-----------\n')
list_of_params = list(params.__dict__.keys())
print_params = ''
for param in list_of_params:
param_value = getattr(params, param)
print_params += f'{param}: {param_value:.2f}'
print_params = print_params[:-1]
f.write(print_params + '\n')
f.write('Best ND: ' + str(best_test_ND) + '\n')
logger.info(print_params)
logger.info(f'Best ND: {best_test_ND}')
f.close()
if batch_n % 200 == 0:
with tf.device('CPU'):
train_sums = list(tf.concat(greedy_seqs.values, axis=0).numpy())
train_inds = list(tf.concat(batch[-1].values, axis=0).numpy().squeeze())
articles = [article[x] for x in train_inds]
gt_summaries = [summary[x] for x in train_inds]
examples_oovs = [oovs[x] for x in train_inds]
scores, summaries, time_step_masks = env.get_rewards(gt_summaries, train_sums, examples_oovs)
save_examples(examples_folder, articles, gt_summaries, summaries, epoch, batch_n, 'train', stage='pretrain')
save_scores(metrics_folder, scores, 'pretrain')
mean_epoch_loss = np.mean(losses)
losses = []
save_loss(metrics_folder, mean_epoch_loss, 'pretrain')
val_losses = []
val_sums = []
val_inds = []
val_iterator = iter(val_dist_dataset)
for val_batch_n in range(1, min(10, batches_per_epoch)):
batch = next(val_iterator)
loss, greedy_seqs, _ = distributed_step(batch, 'val')
val_losses.append(loss)
with tf.device('CPU'):
val_sums += list(tf.concat(greedy_seqs.values, axis=0).numpy())
val_inds += list(tf.concat(batch[-1].values, axis=0).numpy().squeeze())
articles = [val_article[x] for x in val_inds]
def train(config, X_train, Y_train, Train_cs_features, Train_spatial_features,
Train_pattern_features, Train_median_features, Train_q25_features,
Train_q75_features, input_horizon):
'''Y_train = torch.from_numpy(Y_train).float()
X_train = torch.from_numpy(X_train).float()
Train_cs_features = torch.from_numpy(Train_cs_features).float()
Train_spatial_features = torch.from_numpy(Train_spatial_features).float()
Train_pattern_features = torch.from_numpy(Train_pattern_features).float()
Train_median_features = torch.from_numpy(Train_median_features).float()
Train_q25_features = torch.from_numpy(Train_q25_features).float()
Train_q75_features = torch.from_numpy(Train_q75_features).float()'''
Y_train = torch.Tensor(Y_train)
X_train = torch.Tensor(X_train)
#X_train = X_train[:,:,1:] # remove occupancy
Train_cs_features = torch.Tensor(Train_cs_features)
Train_spatial_features = torch.Tensor(Train_spatial_features)
Train_pattern_features = torch.Tensor(Train_pattern_features)
Train_median_features = torch.Tensor(Train_median_features)
Train_q25_features = torch.Tensor(Train_q25_features)
Train_q75_features = torch.Tensor(Train_q75_features)
if len(X_train.shape) == 2: # if 2d make it 3d
X_train = X_train.unsqueeze(
2
) # add 3rd dimesion when not one hot enocded or no additional features
print(X_train.shape)
input_size = X_train.shape[2] # 1 or additional attributes
output_size = 1
#output_size = X_train.shape[2]
hidden_size = int(config['train']['hidden_size'])
embed_size = hidden_size
#Model hyperparameters
lr = float(config['train']['lr'])
dropout = float(config['train']['dropout'])
wd = float(config['train']['wd'])
num_epochs = int(config['train']['num_epochs'])
batch_size = int(config['train']['batch_size'])
num_layers = int(config['train']['num_layers'])
algo = config['train']['algo']
decode = config['model']['decoder']
feat = config.getboolean('data', 'features')
f_name = algo + '_' + str(input_horizon) + '.pth.tar'
if algo == 'seq2seq':
encoder = Encoder(input_size=input_size,
hidden_size=hidden_size,
dropout=dropout,
num_layers=num_layers).to(device)
embedding_cs = Embedding(feat_size=Train_cs_features.shape[1],
embed_size=embed_size)
embedding_spatial = Embedding(
feat_size=Train_spatial_features.shape[1], embed_size=embed_size)
embedding_pattern = Embedding(
feat_size=Train_pattern_features.shape[1], embed_size=embed_size)
embedding_median = Embedding(feat_size=Train_median_features.shape[1],
embed_size=embed_size)
embedding_q25 = Embedding(feat_size=Train_q25_features.shape[1],
embed_size=embed_size)
embedding_q75 = Embedding(feat_size=Train_q75_features.shape[1],
embed_size=embed_size)
embedding = Embedding(feat_size=hidden_size + (3 * embed_size),
embed_size=embed_size)
if decode == 'attention':
decoder = AttnDecoder(
input_size=1,
hidden_size=hidden_size,
output_size=output_size,
input_len=X_train.shape[1],
feat_size_cs=Train_cs_features.shape[1],
feat_size_spatial=Train_spatial_features.shape[1],
dropout=dropout,
num_layers=num_layers).to(device)
if decode == 'decoder':
decoder = Decoder(
input_size=1,
hidden_size=hidden_size,
output_size=output_size,
#feat_size_cs=Train_cs_features.shape[1], feat_size_spatial=Train_spatial_features.shape[1],
dropout=dropout,
num_layers=num_layers).to(device)
'''decoder = Decoder(input_size=X_train.shape[2], hidden_size=hidden_size, output_size=output_size,
feat_size_cs=Train_cs_features.shape[1],
feat_size_spatial=Train_spatial_features.shape[1],
dropout=dropout, num_layers=num_layers).to(device)'''
if decode == 'features':
decoder = Decoder(
input_size=1,
hidden_size=hidden_size,
output_size=output_size,
#feat_size_cs=Train_cs_features.shape[1], feat_size_spatial=Train_spatial_features.shape[1],
dropout=dropout,
num_layers=num_layers).to(device)
'''decoder = Decoder(input_size=1, hidden_size=hidden_size + Train_pattern_features.shape[1], output_size=output_size, feat_size_cs=Train_cs_features.shape[1],
feat_size_spatial=Train_spatial_features.shape[1],
dropout=dropout, num_layers=num_layers).to(device)'''
model = Seq2Seq(encoder, decoder, embedding_cs, embedding_spatial,
embedding_pattern, embedding_median, embedding_q25,
embedding_q75, embedding, config).to(device)
elif algo == 'baseline':
# ouput size = seq length
model = DeepBaseline(input_size=input_size,
hidden_size=hidden_size,
output_size=Y_train.shape[1]).to(device)
criterion = nn.BCELoss()
#criterion = FocalLoss()
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=wd)
pytorch_total_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('Total trainable Params: ', pytorch_total_params)
n_batches_train = int(X_train.shape[0] / batch_size)
print(n_batches_train)
#positive_wt, negative_wt = compute_weights(Y_train)
train_loss = []
for epoch in range(num_epochs):
print('Epoch : ' + str(epoch) + '/' + str(num_epochs))
for b in range(n_batches_train):
b = b * batch_size
#for phase in phases:
#if phase == 'train':
input_batch = X_train[b:b + batch_size, :, :].to(device)
target_label = Y_train[b:b + batch_size, :].to(device) # here
#target_label = Y_train[b: b + batch_size, :, :].to(device) # here
features_cs = Train_cs_features[b:b + batch_size, :].to(device)
features_spatial = Train_spatial_features[b:b +
batch_size, :].to(device)
features_pattern = Train_pattern_features[b:b +
batch_size, :].to(device)
features_median = Train_median_features[b:b +
batch_size, :].to(device)
features_q25 = Train_q25_features[b:b + batch_size, :].to(device)
features_q75 = Train_q75_features[b:b + batch_size, :].to(device)
#positive_wt, negative_wt = compute_weights(target_label)
model.train()
optimizer.zero_grad()
#with torch.set_grad_enabled(phase == 'train'):
if algo == 'seq2seq':
outputs = model(input_batch, target_label, features_cs,
features_spatial, features_pattern,
features_median, features_q25, features_q75,
0.0) # no teacher force
elif algo == 'baseline':
hidden = model.init_hidden(batch_size).to(device)
outputs = model(input_batch, hidden)
#print(target_label.shape, outputs.shape)
#weights = compute_weight_matrix(target_label, positive_wt, negative_wt)
#criterion.weight = weights
loss = criterion(outputs, target_label) # here
print(loss)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1)
optimizer.step()
train_loss.append(loss.item())
checkpoint = {
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}
save_checkpoint(checkpoint, config, filename=f_name)
loss_file = algo + '_' + str(input_horizon) + '.pkl'
#save_loss(config, train_loss, test_loss, loss_file)
save_loss(config, train_loss, loss_file)
print('Finished training')
def run(path_to_net,
label_dir,
nii_dir,
plotter,
batch_size=32,
test_split=0.3,
random_state=666,
epochs=8,
learning_rate=0.0001,
momentum=0.9,
num_folds=5):
"""
Applies training and validation on the network
"""
print('Setting started', flush=True)
nii_filenames = np.asarray(glob.glob(nii_dir + '/*.npy'))
print('Number of files: ', len(nii_filenames), flush=True)
# Creating data indices
dataset_size = len(nii_filenames)
indices = list(range(dataset_size))
test_indices, trainset_indices = utils.get_test_indices(
indices, test_split)
# kfold index generator
for cv_num, (train_idx, val_idx) in enumerate(
utils.get_train_cv_indices(trainset_indices, num_folds,
random_state)):
# take from trainset_indices the kfold generated ones
train_indices = np.asarray(trainset_indices)[np.asarray(train_idx)]
val_indices = np.asarray(trainset_indices)[np.asarray(val_idx)]
print('cv cycle number: ', cv_num, flush=True)
net = Net()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
num_GPU = torch.cuda.device_count()
print('Device: ', device, flush=True)
if num_GPU > 1:
print('Let us use', num_GPU, 'GPUs!', flush=True)
net = nn.DataParallel(net)
net.to(device)
# weigh the loss with the size of classes
# class 0: 3268
# class 1: 60248
weight = torch.tensor([1. / 3268., 1. / 60248.]).to(device)
criterion = nn.CrossEntropyLoss(weight=weight)
optimizer = optim.Adam(net.parameters(), lr=learning_rate)
scheduler = ReduceLROnPlateau(optimizer,
threshold=1e-6,
patience=0,
verbose=True)
fMRI_dataset_train = dataset.fMRIDataset(label_dir,
nii_dir,
train_indices,
transform=dataset.ToTensor())
fMRI_dataset_val = dataset.fMRIDataset(label_dir,
nii_dir,
val_indices,
transform=dataset.ToTensor())
datalengths = {
'train': len(fMRI_dataset_train),
'val': len(fMRI_dataset_val)
}
dataloaders = {
'train': utils.get_dataloader(fMRI_dataset_train, batch_size,
num_GPU),
'val': utils.get_dataloader(fMRI_dataset_val, batch_size, num_GPU)
}
print('Train set length {}, Val set length {}: '.format(
datalengths['train'], datalengths['val']))
# Setup metrics
running_metrics_val = metrics.BinaryClassificationMeter()
running_metrics_train = metrics.BinaryClassificationMeter()
val_loss_meter = metrics.averageLossMeter()
train_loss_meter = metrics.averageLossMeter()
# Track iteration number over epochs for plotter
itr = 0
# Track lowest loss over epochs for saving network
lowest_loss = 100000
for epoch in tqdm(range(epochs), desc='Epochs'):
print('Epoch: ', epoch + 1, flush=True)
print('Phase: train', flush=True)
phase = 'train'
# Set model to training mode
net.train(True)
# Iterate over data.
for i, data in tqdm(enumerate(dataloaders[phase]),
desc='Dataiteration_train'):
train_pred, train_labels, train_loss = train(
data, optimizer, net, criterion, device)
running_metrics_train.update(train_pred, train_labels)
train_loss_meter.update(train_loss, n=1)
if (i + 1) % 10 == 0:
print('Number of Iteration [{}/{}]'.format(
i + 1, int(datalengths[phase] / batch_size)),
flush=True)
itr += 1
score = running_metrics_train.get_scores()
for k, v in score.items():
plotter.plot(k, 'itr', phase, k, itr, v)
print(k, v, flush=True)
print('Loss Train', train_loss_meter.avg, flush=True)
plotter.plot('Loss', 'itr', phase, 'Loss Train', itr,
train_loss_meter.avg)
utils.save_scores(running_metrics_train.get_history(),
phase, cv_num)
utils.save_loss(train_loss_meter.get_history(), phase,
cv_num)
print('Phase: val', flush=True)
phase = 'val'
# Set model to validation mode
net.train(False)
with torch.no_grad():
for i, data in tqdm(enumerate(dataloaders[phase]),
desc='Dataiteration_val'):
val_pred, val_labels, val_loss = val(
data, net, criterion, device)
running_metrics_val.update(val_pred, val_labels)
val_loss_meter.update(val_loss, n=1)
if (i + 1) % 10 == 0:
print('Number of Iteration [{}/{}]'.format(
i + 1, int(datalengths[phase] / batch_size)),
flush=True)
utils.save_scores(running_metrics_val.get_history(), phase,
cv_num)
utils.save_loss(val_loss_meter.get_history(), phase,
cv_num)
if val_loss_meter.avg < lowest_loss:
lowest_loss = val_loss_meter.avg
utils.save_net(path_to_net,
batch_size,
epoch,
cv_num,
train_indices,
val_indices,
test_indices,
net,
optimizer,
criterion,
iter_num=i)
# Plot validation metrics and loss at the end of the val phase
score = running_metrics_val.get_scores()
for k, v in score.items():
plotter.plot(k, 'itr', phase, k, itr, v)
print(k, v, flush=True)
print('Loss Val', val_loss_meter.avg, flush=True)
plotter.plot('Loss', 'itr', phase, 'Loss Val', itr,
val_loss_meter.avg)
print(
'Epoch [{}/{}], Train_loss: {:.4f}, Train_bacc: {:.2f}'.format(
epoch + 1, epochs, train_loss_meter.avg,
running_metrics_train.bacc),
flush=True)
print('Epoch [{}/{}], Val_loss: {:.4f}, Val_bacc: {:.2f}'.format(
epoch + 1, epochs, val_loss_meter.avg,
running_metrics_val.bacc),
flush=True)
# Call the learning rate adjustment function after every epoch
scheduler.step(train_loss_meter.avg)
# Save net after every cross validation cycle
utils.save_net(path_to_net, batch_size, epochs, cv_num, train_indices,
val_indices, test_indices, net, optimizer, criterion)
def train_model(args):
""" Train model """
train_mtx, vocab, train_labels = check_and_load_training_data(args)
# -- configuration --
config = utils.create_baysmm_config(args)
config['vocab_size'] = len(vocab)
config['n_docs'] = train_mtx.shape[0]
config['dtype'] = 'float'
# -- end of configuration --
logging.basicConfig(format='%(asctime)s %(message)s',
datefmt='%d-%m-%Y %H:%M:%S',
filename=config['exp_dir'] + 'training.log', filemode='a',
level=getattr(logging, args.log.upper(), None))
print("Log file:", config['exp_dir'] + 'training.log')
if args.v:
logging.getLogger().addHandler(logging.StreamHandler())
logging.info('PyTorch version: %s', str(torch.__version__))
model = create_model(train_mtx, config, args)
# params = utils.load_params(eng_model_h5_file)
# model.Q.data = params['Q'].to(device=model.device)
# if config['cuda']:
# utils.estimate_approx_num_batches(model, train_mtx)
if args.trn <= config['trn_done']:
logging.info('Found model that is already trained.')
return
config['trn_iters'] = args.trn
optims = create_optimizers(model, config)
# optims['Q'] = None
utils.save_config(config)
dset = utils.SMMDataset(train_mtx, train_labels, len(vocab), 'unsup')
dset.to_device(model.device)
if args.batchwise:
model, loss_iters = batch_wise_training(model, optims, dset,
config, args)
else:
logging.info('Training on {:d} docs.'.format(config['n_docs']))
loss_iters = model.train_me(dset, optims, args.nb)
t_sparsity = utils.t_sparsity(model)
utils.write_info(model.config, "Sparsity in T: {:.2f}".format(t_sparsity))
logging.info("Initial ELBO: {:.1f}".format(-loss_iters[0][0]))
logging.info(" Final ELBO: {:.1f}".format(-loss_iters[-1][0]))
logging.info("Sparsity in T: {:.2f}".format(t_sparsity))
utils.save_model(model)
base = os.path.basename(args.mtx_file).split('.')[0]
sfx = "_T{:d}".format(config['trn_done'])
utils.save_loss(loss_iters, model.config, base, sfx)
def extract_ivector_posteriors(args):
""" Extract posterior distribution of i-vectors using existing model """
# -- configuration --
cfg_f = os.path.dirname(os.path.realpath(args.model_f)) + "/config.json"
config = json.load(open(cfg_f, 'r'))
os.makedirs(config['tmp_dir'] + 'ivecs/', exist_ok=True)
config['xtr_done'] = 0
config['xtr_iters'] = args.xtr
config['nth'] = args.nth
# -- end of configuration
logging.basicConfig(format='%(asctime)s %(message)s',
datefmt='%d-%m-%Y %H:%M:%S',
filename=config['exp_dir'] + 'extraction.log',
level=getattr(logging, args.log.upper(), None))
print("Log file:", config['exp_dir'] + 'extraction.log')
if args.v:
logging.getLogger().addHandler(logging.StreamHandler())
logging.info('PyTorch version: %s, Cuda: %s', str(torch.__version__), str(args.cuda))
mtx_list = get_mtx_list_for_extraction(args.extract_list)
logging.info("Number of files for extraction: %d", len(mtx_list))
params = utils.load_params(args.model_f)
for mtx_file in mtx_list:
data_mtx = sio.mmread(mtx_file).tocsc()
if data_mtx.shape[0] == config['vocab_size']:
data_mtx = data_mtx.T
sbase = os.path.basename(mtx_file).split(".")[0]
mbase = os.path.basename(args.model_f).split(".")[0]
out_file = config['ivecs_dir'] + sbase + "_" + mbase + "_e"
out_file += str(config['xtr_iters']) + ".h5"
if os.path.exists(out_file):
logging.info("i-vector posterior distributions were %s %s",
"already extracted and saved in", out_file)
continue
logging.info('Extracting i-vectors for %d %s', data_mtx.shape[0], 'docs')
# Create model and copy existing parameters
model = BaySMM(params['m'], config, args.cuda)
model.T.data = params['T']
model.T.requires_grad_(False)
model.m.requires_grad_(False)
# Create dataset object
dset = utils.SMMDataset(data_mtx, None, config['vocab_size'], 'unsup')
dset.to_device(model.device)
# Reset i-vector posterior parameters
model.init_var_params(data_mtx.shape[0])
# move model to device (CUDA if available)
model.to_device(model.device)
logging.info('Model Device: %s', str(model.device))
# Create optimizer
if config['optim'] == 'adam':
opt_e = torch.optim.Adam([model.Q], lr=config['eta_q'])
else:
opt_e = torch.optim.Adagrad([model.Q], lr=config['eta_q'])
# extract
loss_iters = model.extract_ivector_posteriors(dset, opt_e, sbase,
args.nb)
sfx = "_" + mbase + "_e{:d}".format(config['xtr_iters'])
utils.save_loss(loss_iters, model.config, "xtr_" + sbase, sfx)
# utils.merge_ivecs(config, sbase, mbase, config['xtr_iters'])
utils.merge_ivecs_v2(config, sbase, mbase, config['xtr_iters'], args.nb)
f'samples{args.samples}'
f'{args.tail}')
os.makedirs(model_dir, exist_ok=True)
utils.save_params(model_dir, vars(args))
print(model_dir)
## Data
trainset, testset, num_classes = L.load_dataset(args.data,
data_dir=args.data_dir)
X_train, y_train = F.get_samples(trainset, args.samples)
X_train, y_train = X_train.to(device), y_train.to(device)
## Architecture
net = L.load_architecture(args.data, args.arch)
net = net.to(device)
## Training
with torch.no_grad():
Z_train = net.init(X_train, y_train)
losses_train = net.get_loss()
X_train, Z_train = F.to_cpu(X_train, Z_train)
## Saving
utils.save_loss(model_dir, 'train', losses_train)
utils.save_ckpt(model_dir, 'model', net)
## Plotting
plot.plot_loss_mcr(model_dir, 'train')
print(model_dir)
def train_GAN(tranfer_GAN, envs, replay_buffer, args):
fixed_noise = Noise(args)
# Loss function
criterion = torch.nn.BCELoss()
# Optimizers
G_optimizer = torch.optim.Adam(tranfer_GAN.G.parameters(), lr=args.G_lr[0])
D_optimizer = torch.optim.Adam(tranfer_GAN.D.parameters(), lr=args.D_lr)
gamma_decay = (args.G_lr[0] - args.G_lr[1]) / (args.gan_num_epochs *
args.gan_num_steps)
print(gamma_decay)
# LR Scheduler
lr_scheduler_G = gammaLR(G_optimizer, gamma=gamma_decay)
# Training GAN
D_avg_losses = []
G_avg_losses = []
replay_buffer.burn_in(envs, args.burn_in_frames)
# Reset Env Before Start
envs.reset()
for epoch in range(args.gan_num_epochs):
D_losses = []
G_losses = []
# Save model Each N Epochs
if (epoch % 10) == 0:
torch.save(tranfer_GAN,
'./models/sonic_VGAN_epoch_' + str(epoch) + '.pt')
for step in range(args.gan_num_steps):
# Feed Memory Replay with Real Sonic Images #
#envs.render()
actions = [envs.action_space.sample()] * envs.num_envs
obs, rewards, dones, info = envs.step(actions)
replay_buffer.append(obs)
# __________________________________________________#
# Train discriminator with real data #
real_x = V(ToTensor(replay_buffer.sample(
args.batch_size))).squeeze() # labels
real_y = V(torch.ones(real_x.size()[0]).cuda())
D_probs = tranfer_GAN.D(real_x + Noise(args).cuda())
D_real_loss = criterion(D_probs, real_y)
# __________________________________________________#
# Get noise to Feed Generator #
fake_image = tranfer_GAN.G(real_x)
fake_y = V(torch.zeros(fake_image.size()[0]).cuda())
last_fake = real_x.data.clone()
last_fake[:, -1:, :, :] = fake_image
D_fake_probs = tranfer_GAN.D(last_fake).squeeze()
D_fake_loss = criterion(D_fake_probs, fake_y)
# Back propagation
D_loss = D_real_loss + D_fake_loss
tranfer_GAN.D.zero_grad()
D_loss.backward()
D_optimizer.step()
# _________________________________________________#
# Train generator
#noise = Noise(args)
real_x = V(ToTensor(replay_buffer.sample(
args.batch_size))).squeeze() # labels
real_y = V(torch.ones(real_x.size()[0]).cuda())
fake_image = tranfer_GAN.G(real_x)
# Get Lp loss from Fake and Real images:
LP_loss = lp_loss(fake_image, real_x[:, -1:, :, :])
last_fake = real_x.data.clone()
last_fake[:, -1:, :, :] = fake_image
D_fake_probs = tranfer_GAN.D(last_fake + Noise(args).cuda())
G_loss = criterion(D_fake_probs, real_y)
# Total loss:
T_G_loss = args.lambda_adv * G_loss + args.lambda_lp * LP_loss
# Back propagation
tranfer_GAN.D.zero_grad()
tranfer_GAN.G.zero_grad()
#G_loss.backward()
T_G_loss.backward()
G_optimizer.step()
lr_scheduler_G.step()
# loss values
D_losses.append(D_loss.data[0])
G_losses.append(G_loss.data[0])
# _________________________________________________#
print('Epoch [%d/%d], Step [%d/%d], D_loss: %.4f, G_loss: %.4f' %
(epoch + 1, args.gan_num_epochs, step + 1,
args.gan_num_steps, D_loss.data[0], G_loss.data[0]))
D_avg_loss = torch.mean(torch.FloatTensor(D_losses))
G_avg_loss = torch.mean(torch.FloatTensor(G_losses))
# avg loss values for plot
D_avg_losses.append(D_avg_loss)
G_avg_losses.append(G_avg_loss)
save_loss(D_avg_losses, G_avg_losses, epoch, save=True)
def train(self):
g_total_loss = []
d_total_loss = []
for epoch in range(args.epochs):
g_epoch_loss = 0
d_epoch_loss = 0
for batch, (imgs, _) in enumerate(self.images_loader):
# Adversarial ground truths
valid = Variable(Tensor(imgs.size(0), 1, 1,
1).fill_(self.real_label),
requires_grad=False)
fake = Variable(Tensor(imgs.size(0), 1, 1,
1).fill_(self.fake_label),
requires_grad=False)
with torch.no_grad():
real_imgs = Variable(imgs.type(Tensor))
# Train Discriminator
self.optimizer_D.zero_grad()
z = torch.randn(imgs.shape[0], args.nz, 1, 1)
# generate fake images from latent vectors
gen_imgs = self.generator(z)
# train on real data
d_real = self.discriminator(
utils.gaussian(real_imgs, 0, args.n_sd))
real_loss = self.adversarial_loss(d_real, args.real)
real_loss.backward()
# train on fake data
d_fake = self.discriminator(
utils.gaussian(gen_imgs.detach(), 0, args.n_sd))
fake_loss = self.adversarial_loss(d_fake, args.fake)
fake_loss.backward()
self.optimizer_D.step()
d_loss = real_loss + fake_loss
# d_loss.backward()
# Train Generator
self.optimizer_G.zero_grad()
gen_imgs = self.generator(z)
d_gen = self.discriminator(
utils.gaussian(gen_imgs, 0, args.n_sd))
g_loss = self.adversarial_loss(d_gen, args.real)
g_loss.backward()
self.optimizer_G.step()
# write losses to files
with torch.no_grad():
utils.save_loss(epoch + 1, batch + 1,
len(self.images_loader), d_loss.item(),
'discriminator')
utils.save_loss(epoch + 1, batch + 1,
len(self.images_loader), g_loss.item(),
'generator')
# Save Losses for plotting later
g_epoch_loss += g_loss.item()
d_epoch_loss += d_loss.item()
batches_done = epoch * len(self.images_loader) + batch
if batches_done % args.sample_interval == 0:
with torch.no_grad():
generated = self.generator(self.fixed_noise)
utils.save_images(generated.detach(),
'generated',
n_row=8,
batches_done=batches_done)
print("Generated samples saved")
g_total_loss.append(g_epoch_loss / len(self.images_loader))
d_total_loss.append(d_epoch_loss / len(self.images_loader))
utils.save_plot(epoch + 1, g_total_loss, "Generator")
utils.save_plot(epoch + 1, d_total_loss, "Discriminator")
print("Plots saved")
# decay learning rate
self.scheduler_D.step()
print('Epoch:', epoch, 'D LR:', self.scheduler_D.get_lr())
self.scheduler_G.step()
print('Epoch:', epoch, 'G LR:', self.scheduler_G.get_lr())
