def train_one_epoch(self, dataloader, optimizer, scheduler, exp_config, type):
model = None
if type == "clean":
model = self.main_model
elif type == "noisy":
model = self.cm_model
batch_iterator = tqdm(dataloader, desc=f"Step {type}")
for step, batch in enumerate(batch_iterator):
model.train()
inputs, labels = batch_to_inputs(batch, exp_config)
if type == "clean":
logits = model(**inputs)[0]
loss = compute_loss(logits, labels, is_log_prob=False)
elif type == "noisy":
log_probs = model(inputs)
loss = compute_loss(log_probs, labels, is_log_prob=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), exp_config["max_grad_norm"])
optimizer.step()
scheduler.step()
model.zero_grad()
def _train_d(self, real_batch, noise_batch):
self.optim_d.zero_grad()
models.toggle_grad(self.G, False)
models.toggle_grad(self.D, True)
# train real
if len(real_batch) == 2:
real_batch_x, real_batch_y = real_batch
real_batch_x = real_batch_x.to(self.device)
real_batch_y = real_batch_y.to(self.device)
else:
real_batch_x = real_batch.to(self.device)
real_batch_y = None
pred_real = self.D(real_batch_x, real_batch_y)
loss_real = compute_loss(self.loss, pred_real, 1)
loss_real.backward()
# train fake
fake_batch = self.G(noise_batch, real_batch_y).data
pred_fake = self.D(fake_batch, real_batch_y)
loss_fake = compute_loss(self.loss, pred_fake, 0)
loss_fake.backward()
# grad pen
grad_pen = cal_grad_pen(self.D,
real_batch=real_batch_x,
fake_batch=fake_batch,
gp_weight=self.args['training']['gp_weight'],
gp_inter=self.args['training']['gp_inter'],
gp_center=self.args['training']['gp_center'],
label_batch=real_batch_y)
grad_pen.backward()
# update params
self.optim_d.step()
loss_d = loss_real + loss_fake + grad_pen
return loss_d.item()
def _train_g(self, noise_batch, label_batch):
self.optim_g.zero_grad()
models.toggle_grad(self.G, True)
models.toggle_grad(self.D, False)
fake_batch = self.G(noise_batch, label_batch)
pred_fake = self.D(fake_batch, label_batch)
loss_fake = compute_loss(self.loss, pred_fake, 1)
loss_fake.backward()
# update params
self.optim_g.step()
return loss_fake.item()
def train_one_epoch(self, dataloader, optimizer, scheduler, exp_config,
type):
self.main_model.train()
batch_iterator = tqdm(dataloader, desc=f"Step {type}")
for step, batch in enumerate(batch_iterator):
inputs, labels = batch_to_inputs(batch, exp_config)
clean_outputs = self.main_model(**inputs)
logits = clean_outputs[
0] # if we do not provide labels to the BERT model, the first output is the logits
loss = compute_loss(logits, labels)
loss.backward()
torch.nn.utils.clip_grad_norm_(self.main_model.parameters(),
exp_config["max_grad_norm"])
optimizer.step()
scheduler.step()
self.main_model.zero_grad()
num_tests = 5 # Number of times to test for test accuracy
max_checkpoints_to_keep = 3
save_dir = "data/checkpoints"
train_vars = 'models/fc8-pets/weights:0,models/fc8-pets/biases:0'
# Export
export_dir = "/tmp/export/"
export_name = "pet-model" export_version = 2
images, labels = datasets.input_pipeline(dataset_dir, batch_size, is_training=True)
test_images, test_labels = datasets.input_pipeline(dataset_dir, batch_size, is_training=False)
with tf.variable_scope("models") as scope:
logits = nets.inference(images, is_training=True)
scope.reuse_variables()
test_logits = nets.inference(test_images, is_training=False)
total_loss = models.compute_loss(logits, labels)
train_accuracy = models.compute_accuracy(logits, labels)
test_accuracy = models.compute_accuracy(test_logits, test_labels)
global_step = tf.Variable(0, trainable=False)
learning_rate = models.get_learning_rate(global_step, initial_learning_rate, decay_steps, decay_rate)
train_op = models.train(total_loss, learning_rate, global_step, train_vars)
saver = tf.train.Saver(max_to_keep=max_checkpoints_to_keep)
checkpoints_dir = os.path.join(save_dir, datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
if not os.path.exists(save_dir):
os.mkdir(save_dir)
if not os.path.exists(checkpoints_dir):
os.mkdir(checkpoints_dir)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
model = Sequential()
first_layer = Layer(4, "sigmoid")
model.add(first_layer)
second_layer = Layer(5, "sigmoid")
model.add(second_layer)
third_layer = Layer(4, "softmax")
model.add(third_layer)
model.compile()
loss = 0
for i in range(EPOCHS):
inpt = (uniform(-1, 1), uniform(-1, 1), uniform(-1, 1), uniform(-1, 1))
expected_output = [1 if n == max(inpt) else 0 for n in inpt]
output = model.run(inpt)
loss += compute_loss(output, expected_output)
if loss < min_loss:
best_model = model
min_loss = loss
print("Loss is: " + str(loss))
# inp = (0.2, 0.2, 0.1, 0.3)
# r = best_model.run(inp)
# out = r.index(max(r)) + 1
# expected_output = [1 if n == max(inp) else 0 for n in inp]
# print("TEST: " + str(out) + " -- Expected: " + str(expected_output))
# inp = (0.4, 0.2, -0.1, 0.3)
# r = best_model.run(inp)
# out1 = r.index(max(r)) + 1
# expected_output = [1 if n == max(inp) else 0 for n in inp]
train_input_queue = data_flow_ops.FIFOQueue(capacity=10000,
dtypes=[tf.string, tf.int64],
shapes=[(num_frames, ), ()])
train_enqueue_op = train_input_queue.enqueue_many(
[image_paths_placeholder, labels_placeholder])
frames_batch, labels_batch = input_pipeline(train_input_queue,
batch_size=batch_size,
image_size=image_size)
with tf.variable_scope("models") as scope:
logits, _ = nets.inference(frames_batch, is_training=True)
total_loss, cross_entropy_loss, reg_loss = models.compute_loss(
logits, labels_batch)
train_accuracy = models.compute_accuracy(logits, labels_batch)
global_step = tf.Variable(0, trainable=False)
learning_rate = models.get_learning_rate(global_step, initial_learning_rate,
decay_steps, decay_rate)
train_op = models.train(total_loss, learning_rate, global_step)
tf.summary.scalar("learning_rate", learning_rate)
tf.summary.scalar("train/accuracy", train_accuracy)
tf.summary.scalar("train/total_loss", total_loss)
tf.summary.scalar("train/cross_entropy_loss", cross_entropy_loss)
tf.summary.scalar("train/regularization_loss", reg_loss)
summary_op = tf.summary.merge_all()
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
total_gradients = []
frames_batch_split = tf.split(frames_batch, num_gpu)
labels_batch_split = tf.split(labels_batch, num_gpu)
print("frames_batch_split", frames_batch_split)
for i in range(num_gpu):
with tf.device('/gpu:%d' % i):
with tf.variable_scope(tf.get_variable_scope(), reuse=(i > 0)):
print("Setup gpu:%d" % i)
print("frames_batch_split_i", i, frames_batch_split[i])
logits_split, _ = nets.inference(frames_batch_split[i], is_training=True)
labels_split = labels_batch_split[i]
total_loss, cross_entropy_loss, reg_loss = models.compute_loss(logits_split, labels_split)
grads = optimizer.compute_gradients(total_loss)
total_gradients.append(grads)
tf.get_variable_scope().reuse_variables()
with tf.device('/cpu:0'):
gradients = models.average_gradients(total_gradients)
train_op = optimizer.apply_gradients(gradients, global_step)
train_accuracy = models.compute_accuracy(logits_split, labels_split)
tf.summary.scalar("learning_rate", learning_rate)
tf.summary.scalar("train/accuracy", train_accuracy)