def train_epoch(model,
dataset,
device,
*,
reg_weight,
learning_rate,
batch_size,
loss_print_freq=500,
max_batches=None):
loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
with tqdm(enumerate(loader), total=len(dataset) / batch_size,
desc='Batch') as pipe:
for batch_idx, batch in pipe:
for k, v in batch.items():
batch[k] = v.to(device)
y, (v_a, _), (v_b, m_b) = model(batch['frame_a'], batch['frame_b'],
batch['amplification_f'])
v_c, m_c = model.encoder(batch['frame_perturbed'])
loss = F.l1_loss(y, batch['frame_amplified']) + \
reg_weight * F.l1_loss(v_a, v_b) + \
reg_weight * F.l1_loss(v_a, v_c) + \
reg_weight * F.l1_loss(m_b, m_c)
if batch_idx % loss_print_freq == 0:
pipe.write("Batch %d loss %.2f" % (batch_idx + 1, loss.item()))
optimizer.zero_grad()
loss.backward()
optimizer.step()
if max_batches is not None and batch_idx + 1 >= max_batches:
break
def calculate_model_losses(args, skip_pixel_loss, img, img_pred, bbox,
bbox_pred):
total_loss = torch.zeros(1).to(img)
losses = {}
l1_pixel_weight = args.l1_pixel_loss_weight
if skip_pixel_loss:
l1_pixel_weight = 0
l1_pixel_loss = F.l1_loss(img_pred, img)
total_loss = add_loss(total_loss, l1_pixel_loss, losses, 'L1_pixel_loss',
l1_pixel_weight)
loss_bbox = F.mse_loss(bbox_pred, bbox)
total_loss = add_loss(total_loss, loss_bbox, losses, 'bbox_pred',
args.bbox_pred_loss_weight)
return total_loss, losses
def compute_critic_loss(self, trajectory):
states = [transition[0] for transition in trajectory]
rewards = [transition[2] for transition in trajectory]
dones = [transition[3] for transition in trajectory]
discounted_rewards = []
cumulative_reward = 0
for step in reversed(range(len(rewards))):
cumulative_reward = rewards[step] + self.gamma * cumulative_reward * (1 - int(dones[step]))
discounted_rewards.insert(0, cumulative_reward)
discounted_rewards = torch.FloatTensor(discounted_rewards).to(self.device)
target_values = torch.FloatTensor(rewards).view(-1, 1).to(self.device) + discounted_rewards.view(-1, 1)
states = torch.cat(states).to(self.device)
_, actual_values = self.model.forward(states)
critic_loss = F.l1_loss(actual_values, target_values.view(-1, 1))
advantage = target_values - actual_values
return critic_loss, advantage.detach()
# e.g {'x': 'a a e b e', 'y': [2, 1, 0, 0, 2]}
writer = SummaryWriter("runs/test-6")
optimizer = optim.Adam(model.parameters(), lr=0.01)
epochs = 200
log_every_n = 1000
for epoch in range(epochs):
for (i, batch) in enumerate(training_data):
batch_x = [b["x"] for b in batch]
batch_y = torch.stack([torch.tensor(b["y"]) for b in batch])
optimizer.zero_grad()
y_hat = model(batch_x)
loss = F.l1_loss(y_hat, batch_y)
writer.add_scalar("loss", loss, epoch * n_training_samples + i)
if i % log_every_n == 0:
print(loss)
loss.backward()
optimizer.step()
writer.close()
torch.save({"model_state_dict": model.state_dict()}, "./saves/good2.pt")
# Things to play around with:
#
def compute_loss_with_l1(loss, model, coef=2e-5):
l1_loss = sum(
F.l1_loss(p, torch.zeros_like(p)) for p in model.parameters()
if p.requires_grad)
return loss + coef * l1_loss
