def train(nbatches, npred):
model.train()
total_loss_i, total_loss_s, total_loss_p = 0, 0, 0
for i in range(nbatches):
optimizer.zero_grad()
inputs, actions, targets, _, _ = dataloader.get_batch_fm(
'train', npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pred, loss_p = model(inputs, actions, targets, z_dropout=opt.z_dropout)
loss_p = loss_p[0]
loss_i, loss_s = compute_loss(targets, pred)
loss = loss_i + loss_s + opt.beta * loss_p
# VAEs get NaN loss sometimes, so check for it
if not math.isnan(loss.item()):
loss.backward(retain_graph=False)
if not math.isnan(utils.grad_norm(model).item()):
torch.nn.utils.clip_grad_norm_(model.parameters(),
opt.grad_clip)
optimizer.step()
total_loss_i += loss_i.item()
total_loss_s += loss_s.item()
total_loss_p += loss_p.item()
del inputs, actions, targets
total_loss_i /= nbatches
total_loss_s /= nbatches
total_loss_p /= nbatches
return total_loss_i, total_loss_s, total_loss_p
def train(nbatches, npred):
model.train()
total_loss = 0
for i in range(nbatches):
optimizer.zero_grad()
inputs, actions, targets, _, _ = dataloader.get_batch_fm(
'train', npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pred, _ = model(inputs, actions, targets, z_dropout=0)
pred_cost = cost(
pred[0].view(opt.batch_size * opt.npred, 1, 3, opt.height,
opt.width), pred[1].view(opt.batch_size * opt.npred,
1, 4))
loss = F.mse_loss(pred_cost.view(opt.batch_size, opt.npred, 2),
targets[2])
if not math.isnan(loss.item()):
loss.backward(retain_graph=False)
if not math.isnan(utils.grad_norm(model).item()):
torch.nn.utils.clip_grad_norm_(model.parameters(),
opt.grad_clip)
optimizer.step()
total_loss += loss.item()
del inputs, actions, targets
total_loss /= nbatches
return total_loss
def test(nbatches, npred):
gamma_mask = torch.Tensor([opt.gamma**t for t in range(npred)]).view(1, -1).cuda()
model.eval()
total_loss_i, total_loss_s, total_loss_c, total_loss_policy, total_loss_p, n_updates = 0, 0, 0, 0, 0, 0
for i in range(nbatches):
inputs, actions, targets, _, _ = dataloader.get_batch_fm('test', npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pred, pred_actions = planning.train_policy_net_mbil(model, inputs, targets, targetprop = opt.targetprop, dropout=0.0, model_type = model_type)
loss_i, loss_s, loss_c_, loss_p = compute_loss(targets, pred)
loss_policy = loss_i + loss_s
if opt.loss_c == 1:
loss_policy += loss_c_
if not math.isnan(loss_policy.item()):
total_loss_i += loss_i.item()
total_loss_s += loss_s.item()
total_loss_p += loss_p.item()
total_loss_policy += loss_policy.item()
n_updates += 1
del inputs, actions, targets, pred
total_loss_i /= n_updates
total_loss_s /= n_updates
total_loss_c /= n_updates
total_loss_policy /= n_updates
total_loss_p /= n_updates
return total_loss_i, total_loss_s, total_loss_c, total_loss_policy, total_loss_p
def compute_pz(nbatches):
model.p_z = []
for j in range(nbatches):
print('[estimating z distribution: {:2.1%}]'.format(
float(j) / nbatches),
end="\r")
inputs, actions, targets = dataloader.get_batch_fm(
'train', opt.npred, (opt.cuda == 1))
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = utils.Variable(actions)
pred, loss_kl = model(inputs, actions, targets, save_z=True)
del inputs, actions, targets
def train(nbatches, npred):
gamma_mask = torch.Tensor([opt.gamma**t
for t in range(npred)]).view(1, -1).cuda()
model.eval()
model.policy_net.train()
total_loss_i, total_loss_s, total_loss_c, total_loss_policy, total_loss_p, n_updates = 0, 0, 0, 0, 0, 0
for i in range(nbatches):
optimizer.zero_grad()
inputs, actions, targets, _, _ = dataloader.get_batch_fm(
'train', npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pred, _ = planning.train_policy_net_mper(model,
inputs,
targets,
dropout=opt.p_dropout,
model_type=model_type)
loss_i, loss_s, loss_c_, loss_p = compute_loss(targets, pred)
# proximity_cost, lane_cost = pred[2][:, :, 0], pred[2][:, :, 1]
# proximity_cost = proximity_cost * Variable(gamma_mask)
# lane_cost = lane_cost * Variable(gamma_mask)
# loss_c = proximity_cost.mean() + opt.lambda_lane * lane_cost.mean()
loss_policy = loss_i + loss_s + opt.lambda_h * loss_p
if opt.loss_c == 1:
loss_policy += loss_c_
if not math.isnan(loss_policy.item()):
loss_policy.backward()
torch.nn.utils.clip_grad_norm(model.policy_net.parameters(),
opt.grad_clip)
optimizer.step()
total_loss_i += loss_i.item()
total_loss_s += loss_s.item()
total_loss_p += loss_p.item()
total_loss_policy += loss_policy.item()
n_updates += 1
else:
print('warning, NaN')
del inputs, actions, targets, pred
total_loss_i /= n_updates
total_loss_s /= n_updates
total_loss_c /= n_updates
total_loss_policy /= n_updates
total_loss_p /= n_updates
return total_loss_i, total_loss_s, total_loss_c, total_loss_policy, total_loss_p
def test(nbatches):
policy.eval()
total_loss, nb = 0, 0
for i in range(nbatches):
inputs, actions, targets, _, _ = dataloader.get_batch_fm('valid')
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pi, mu, sigma, _ = policy(inputs[0], inputs[1])
loss = utils.mdn_loss_fn(pi, sigma, mu,
actions.view(opt.batch_size, -1))
if not math.isnan(loss.item()):
total_loss += loss.item()
nb += 1
else:
print('warning, NaN')
return total_loss / nb
def test(nbatches, npred):
model.train()
model.policy_net.train()
total_loss_c, total_loss_u, total_loss_l, total_loss_a, n_updates = 0, 0, 0, 0, 0
total_loss_policy = 0
for i in range(nbatches):
inputs, actions, targets, ids, car_sizes = dataloader.get_batch_fm(
'valid', npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
pred, actions, _ = planning.train_policy_net_mpur(model,
inputs,
targets,
car_sizes,
n_models=10,
lrt_z=1.0,
n_updates_z=0)
loss_c = pred[2]
loss_l = pred[3]
loss_u = pred[4]
loss_a = actions.norm(2, 2).pow(2).mean()
loss_policy = loss_c + opt.u_reg * loss_u + opt.lambda_l * loss_l + opt.lambda_a * loss_a
if not math.isnan(loss_policy.item()):
total_loss_c += loss_c.item()
total_loss_u += loss_u.item()
total_loss_a += loss_a.item()
total_loss_l += loss_l.item()
total_loss_policy += loss_policy.item()
n_updates += 1
else:
print('warning, NaN')
del inputs, actions, targets, pred
total_loss_c /= n_updates
total_loss_l /= n_updates
total_loss_u /= n_updates
total_loss_a /= n_updates
total_loss_policy /= n_updates
return total_loss_c, total_loss_l, total_loss_u, total_loss_a, total_loss_policy
def test(nbatches, npred):
model.train()
total_loss = 0
for i in range(nbatches):
inputs, actions, targets, _, _ = dataloader.get_batch_fm(
'valid', npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pred, _ = model(inputs, actions, targets, z_dropout=0)
pred_cost = cost(
pred[0].view(opt.batch_size * opt.npred, 1, 3, opt.height,
opt.width), pred[1].view(opt.batch_size * opt.npred,
1, 4))
loss = F.mse_loss(pred_cost.view(opt.batch_size, opt.npred, 2),
targets[2])
if not math.isnan(loss.item()):
total_loss += loss.item()
del inputs, actions, targets
total_loss /= nbatches
return total_loss
def _request_data(self, loader, volatile=False):
try:
batch = next(self.data_iter)
except StopIteration:
self.data_iter = None
return None
batch = utils.make_variables(batch, volatile=volatile)
if self.cuda != '':
batch = utils.cudaify(batch)
return batch
def test(nbatches):
model.eval()
total_loss_i, total_loss_s, total_loss_p = 0, 0, 0
for i in range(nbatches):
inputs, actions, targets, _, _ = dataloader.get_batch_fm('valid')
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pred, loss_p = model(inputs, actions, targets, z_dropout=opt.z_dropout)
loss_p = loss_p[0]
loss_i, loss_s = compute_loss(targets, pred)
loss = loss_i + loss_s + opt.beta * loss_p
total_loss_i += loss_i.item()
total_loss_s += loss_s.item()
total_loss_p += loss_p.item()
del inputs, actions, targets
total_loss_i /= nbatches
total_loss_s /= nbatches
total_loss_p /= nbatches
return total_loss_i, total_loss_s, total_loss_p
def train(nbatches):
policy.train()
total_loss, nb = 0, 0
for i in range(nbatches):
optimizer.zero_grad()
inputs, actions, targets, _, _ = dataloader.get_batch_fm('train')
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
pi, mu, sigma, _ = policy(inputs[0], inputs[1])
loss = utils.mdn_loss_fn(pi, sigma, mu,
actions.view(opt.batch_size, -1))
if not math.isnan(loss.item()):
loss.backward()
if opt.grad_clip != -1:
torch.nn.utils.clip_grad_norm_(policy.parameters(),
opt.grad_clip)
optimizer.step()
total_loss += loss.item()
nb += 1
else:
print('warning, NaN')
return total_loss / nb
def compute_seg_score(dataset, seg_score):
from torch import Tensor
from utils import cudaify, make_variables
values = OrderedDict()
for data in dataset:
pred = maybe_convert_to_magnitude(data[PRED_KEY])
pred = Tensor(pred).unsqueeze(0)
target = Tensor(data[TARGET_LABEL_KEY]).unsqueeze(0)
pred, target = make_variables((pred, target), volatile=True)
if seg_score.cuda != '':
pred, target = cudaify((pred, target))
index_key = _get_index_key(data)
value = seg_score(pred, target)
values[index_key] = value
return pd.Series(values)
def train(nbatches, npred):
model.train()
model.policy_net.train()
total_loss_c, total_loss_u, total_loss_l, total_loss_a, n_updates, grad_norm = 0, 0, 0, 0, 0, 0
total_loss_policy = 0
for j in range(nbatches):
optimizer.zero_grad()
inputs, actions, targets, ids, car_sizes = dataloader.get_batch_fm(
'train', npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
pred, actions, pred_adv = planning.train_policy_net_mpur(
model,
inputs,
targets,
car_sizes,
n_models=10,
lrt_z=opt.lrt_z,
n_updates_z=opt.z_updates,
infer_z=(opt.infer_z == 1))
loss_c = pred[2] # proximity cost
loss_l = pred[3] # lane cost
loss_u = pred[4] # uncertainty cost
loss_a = actions.norm(2, 2).pow(2).mean() # action regularisation
loss_policy = loss_c + opt.u_reg * loss_u + opt.lambda_l * loss_l + opt.lambda_a * loss_a
if not math.isnan(loss_policy.item()):
loss_policy.backward() # back-propagation through time!
grad_norm += utils.grad_norm(model.policy_net).item()
torch.nn.utils.clip_grad_norm_(model.policy_net.parameters(),
opt.grad_clip)
optimizer.step()
total_loss_c += loss_c.item() # proximity cost
total_loss_u += loss_u.item() # uncertainty (reg.)
total_loss_a += loss_a.item() # action (reg.)
total_loss_l += loss_l.item() # lane cost
total_loss_policy += loss_policy.item() # overall total cost
n_updates += 1
else:
print('warning, NaN') # Oh no... Something got quite f****d up!
pdb.set_trace()
if j == 0 and opt.save_movies:
# save videos of normal and adversarial scenarios
for b in range(opt.batch_size):
utils.save_movie(opt.model_file + f'.mov/sampled/mov{b}',
pred[0][b], pred[1][b], None, actions[b])
if pred_adv[0] is not None:
utils.save_movie(
opt.model_file + f'.mov/adversarial/mov{b}',
pred_adv[0][b], pred_adv[1][b], None, actions[b])
del inputs, actions, targets, pred
total_loss_c /= n_updates
total_loss_u /= n_updates
total_loss_a /= n_updates
total_loss_l /= n_updates
total_loss_policy /= n_updates
print(f'[avg grad norm: {grad_norm / n_updates}]')
return total_loss_c, total_loss_l, total_loss_u, total_loss_a, total_loss_policy
loss_c = F.mse_loss(pred_costs, target_costs, reduce=r)
return loss_i, loss_s, loss_c
dataloader.random.seed(12345)
model.train()
ui_truth, ui_perm, ui_turn = [], [], []
us_truth, us_perm, us_turn = [], [], []
for i in range(opt.n_batches):
print(i)
torch.cuda.empty_cache()
inputs, actions, targets, ids, car_sizes = dataloader.get_batch_fm(
'train', opt.npred)
inputs = utils.make_variables(inputs)
targets = utils.make_variables(targets)
actions = Variable(actions)
input_images, input_states = inputs[0], inputs[1]
u_i, u_s, u_c, _, _, _, _ = planning.compute_uncertainty_batch(
model, input_images, input_states, actions, targets, car_sizes)
pred, loss_p = model(inputs, actions, targets, z_dropout=0)
ui_truth.append(u_i)
us_truth.append(u_s)
actions_perm = actions[torch.randperm(4)]
u_i, u_s, u_c, _, _, _, _ = planning.compute_uncertainty_batch(
model, input_images, input_states, actions_perm, targets, car_sizes)
pred, loss_p = model(inputs, actions_perm, targets, z_dropout=0)
ui_perm.append(u_i)
parser.add_argument('-dataset', type=str, default='i80')
parser.add_argument('-batch_size', type=int, default=32)
parser.add_argument('-debug', action='store_true')
parser.add_argument('-seed', type=int, default=9999)
parser.add_argument('-ncond', type=int, default=20)
parser.add_argument('-npred', type=int, default=200)
opt = parser.parse_args()
random.seed(opt.seed)
np.random.seed(opt.seed)
torch.manual_seed(opt.seed)
dataloader = DataLoader(None, opt, opt.dataset)
inputs, actions, targets, ids, car_sizes = dataloader.get_batch_fm('train')
inputs = utils.make_variables(inputs)
images, states = inputs[0], inputs[1]
_, mask = utils.proximity_cost(images,
states,
car_size=car_sizes,
unnormalize=True,
s_mean=dataloader.s_mean,
s_std=dataloader.s_std)
states2 = states[:, -1].clone()
states2 = states2 * (1e-8 + dataloader.s_std.view(1, 4)).cuda()
states2 = states2 + dataloader.s_mean.view(1, 4).cuda()
plt.imshow(images[1][-1].permute(1, 2, 0))
plt.axis('off')
frame1 = plt.gca()
