class ConvModel(Model):
def __init__(self, config):
super().__init__(config)
if config.type == 'conv':
# self.dynamics = ConvTransitionModel2().cuda()
self.dynamics = ConvTransitionModel2_2().cuda() # uses corrected action
self.get_dataset_sample = self.get_dataset_sample_no_speed
self.criterion = F.mse_loss
elif config.type == 'conv_speed':
self.dynamics = ConvTransitionModel3().cuda()
self.get_dataset_sample = self.get_dataset_sample_with_speed
self.criterion = F.mse_loss
elif config.type == 'class':
# self.dynamics = ClassificationModel().cuda()
# self.dynamics = ClassificationModel2().cuda() # uses corrected phase action
self.dynamics = ClassificationModel3().cuda() # uses limited phase history
self.get_dataset_sample = self.get_dataset_sample_for_classification
self.criterion = torch.nn.BCELoss()
elif config.type == 'latent_fc':
self.dynamics = LatentFCTransitionModel().cuda()
self.get_dataset_sample = self.get_dataset_sample_for_latent_fc
self.criterion = F.mse_loss
else:
raise NotImplementedError
self.optim = torch.optim.Adam(self.dynamics.parameters())
self.earlystopping = EarlyStopping(patience=self._c.early_stop_patience)
self.set_epoch_length()
self.writer = SummaryWriter(log_dir=config.logdir, purge_step=0)
def preprocess(self,):
pass
# def get_sample(self):
# return get_dataset_sample(self._dataset)
# danijar style get sample
# yield method
# choose any episode
# why yield episode and not sample.
# while true
# for files in directory:
# if not in cache add to cache
# for i in random set of cache:
# length limitation?
# yield i episode
# while true
# check for files in dir, add new files to cache
# for i in train_steps number of episodes (sampled from episode cache):
# yield a sample of given length
pass
def get_dataset_sample_no_speed(self, dataset):
s = dataset if isinstance(dataset, dict) else next(dataset)
sample = {}
sample['phases'] = torch.Tensor(s['phases'][:, 0, 0, 3, :, 0].numpy()).cuda()
sample['y'] = self.preprocess(torch.Tensor(s['x'][:, 1, :, :, :, 0].numpy()))
sample['v'] = self.preprocess(torch.Tensor(s['x'][:, 0, :, :, :, 1].numpy())) + 0.5
sample['x'] = self.preprocess(torch.Tensor(s['x'][:, 0, :, :, :, 0].numpy()))
sample['action'] = torch.Tensor(s['corrected_action'][:, :1].numpy()).cuda()
## not needed for now.
sample['reward'] = s['reward'].numpy()
# sample['action'] = s['action'].numpy()
return sample
def get_dataset_sample_with_speed(self, dataset):
s = dataset if isinstance(dataset, dict) else next(dataset)
sample = {}
sample['phases'] = self.preprocess(torch.Tensor(s['phases'][:, 0, 0, 3, :, 0].numpy()))
sample['y'] = self.preprocess(torch.Tensor(s['x'][:, 1, :, :, :, :].numpy())).permute(0, 4, 2, 3, 1).squeeze(-1).contiguous()
sample['v'] = self.preprocess(torch.Tensor(s['x'][:, 0, :, :, :, 1].numpy())) + 0.5
sample['x'] = self.preprocess(torch.Tensor(s['x'][:, 0, :, :, :, :].numpy())).permute(0, 4, 2, 3, 1).squeeze(-1).contiguous()
sample['x'][:, 1] = sample['x'][:, 1] + 0.5
sample['y'][:, 1] = sample['y'][:, 1] + 0.5
## not needed for now.
sample['reward'] = s['reward'].numpy()
# sample['action'] = s['action'].numpy()
return sample
def get_dataset_sample_for_classification(self, dataset):
s = dataset if isinstance(dataset, dict) else next(dataset)
sample = {}
sample['phases'] = torch.Tensor(s['phases'][:, 0, 0, 3, :, 0].numpy()).cuda()
sample['y'] = torch.Tensor(s['x'][:, 1, :, :, :, 0].numpy()).cuda()
sample['v'] = torch.Tensor(s['x'][:, 0, :, :, :, 1].numpy()).cuda()
sample['x'] = torch.Tensor(s['x'][:, 0, :, :, :, 0].numpy()).cuda()
sample['action'] = torch.Tensor(s['corrected_action'][:, :1].numpy()).cuda()
sample['phase_action'] = torch.Tensor(s['corrected_p_action'][:, 0].numpy()).cuda()
# classification model only works on the last lane
sample['x'] = sample['x'][:, 0, -1]
sample['y'] = sample['y'][:, 0, -1]
## not needed for now.
sample['reward'] = s['reward'].numpy()
return sample
def get_dataset_sample_for_classification_kstep(self, dataset):
# to see accuracy of k-step predictions
# need formatted samples of higher batch_length
s = dataset if isinstance(dataset, dict) else next(dataset)
sample = {}
sample['phases'] = torch.Tensor(s['phases'][:, :, 0, 3, :, 0].numpy()).cuda()
sample['x'] = torch.Tensor(s['x'][:, :, :, :, :, 0].numpy()).cuda()
sample['action'] = torch.Tensor(s['corrected_action'][:, :].numpy()).cuda()
sample['phase_action'] = torch.Tensor(s['corrected_p_action'][:, :].numpy()).cuda()
# classification model only works on the last lane
sample['x'] = sample['x'][:, :, 0, -1]
sample['reward'] = s['reward'].numpy()
return sample
def get_dataset_sample_for_latent_fc(self, dataset):
s = dataset if isinstance(dataset, dict) else next(dataset)
sample = {}
sample['phases'] = torch.Tensor(s['phases'][:, 0, 0, 3, :, 0].numpy()).cuda()
sample['action'] = torch.Tensor(s['corrected_action'][:, :1].numpy()).cuda()
sample['reward'] = s['reward'].numpy()
mu = (torch.Tensor(s['mu'].numpy())).cuda()
logvar = (torch.Tensor(s['logvar'].numpy())).cuda()
latent = reparameterize(mu, logvar)
sample['x'] = latent[:, 0]
sample['y'] = latent[:, 1]
return sample
def preprocess(self, x):
x = x - 0.5
return x.cuda()
def set_epoch_length(self):
"""
These many number of batches when sampled from the dataset would lead to 1 epoch.
"""
num_episodes = len(self.train_eps)
episode_length = 500
batch_length = self._c.batch_length
batch_size = self._c.batch_size
self.epoch_length = ceil(num_episodes * (episode_length - (batch_length - 1)) / batch_size)
test_num_episodes = len(self.test_eps)
self.test_epoch_length = ceil(test_num_episodes * (episode_length - (batch_length - 1)) / batch_size)
def batch_update_model(self):
# calculate loss
# loss.backward()
# optim.step()
sample = self.get_sample()
loss = self._loss(sample)
loss.backward()
self.optim.step()
def train(self):
cur_best = None
for epoch in range(self._c.epochs):
self.train_dynamics(epoch)
test_loss = self.test(epoch)
# scheduler.step(test_loss)
self.earlystopping.step(test_loss)
self.writer.file_writer.flush()
# checkpointing
best_filename = self._c.logdir / 'best.tar'
filename = self._c.logdir / f'checkpoint_{epoch}.tar'
is_best = not cur_best or test_loss < cur_best
if is_best:
cur_best = test_loss
if is_best or (epoch % 10 == 0):
checkpoint = {
'epoch': epoch,
'state_dict': self.dynamics.state_dict(),
'precision': test_loss,
'optimizer': self.optim.state_dict(),
'earlystopping': self.earlystopping.state_dict(),
# 'scheduler': scheduler.state_dict(),
}
save_checkpoint(checkpoint, is_best, filename, best_filename)
if self.earlystopping.stop:
print("End of Training because of early stopping at epoch {}".format(epoch))
def train_dynamics(self, epoch):
print('=======================> epoch:', epoch)
self.dynamics.train()
train_loss = 0
t1 = time.time()
for u in range(self.epoch_length):
s = self.get_dataset_sample(self._dataset)
self.optim.zero_grad()
y_pred = self.dynamics(s)
loss = self.criterion(y_pred, s['y'])
loss.backward()
train_loss += loss
self.optim.step()
if (u % int(self.epoch_length/min(self.epoch_length, 5)) == 0):
t2 = time.time()
print(u, round(t2-t1, 2), '{:.10f}'.format(loss.item() / self._c.batch_size))
norm_train_loss = (train_loss / (self.epoch_length * self._c.batch_size)).item()
self.writer.add_scalar('train/loss', norm_train_loss, epoch)
print('====> Epoch: {} Average loss: {:.10f}'.format(epoch, norm_train_loss))
def test(self, epoch):
self.dynamics.eval()
test_loss = 0
for u in range(self.test_epoch_length):
s = self.get_dataset_sample(self._test_dataset)
y_pred = self.dynamics(s)
test_loss += F.mse_loss(y_pred, s['y'])
norm_test_loss = (test_loss / (self.test_epoch_length * self._c.batch_size)).item()
self.writer.add_scalar('test/loss', norm_test_loss, epoch)
print('====> Test set loss: {:.10f}'.format(norm_test_loss))
print()
return norm_test_loss
def save(self):
raise NotImplementedError
def load(self):
raise NotImplementedError
def _loss(self):
raise NotImplementedError
def create_reconstructions(self):
pass
def train_model(self):
args = self.args
model = self.model.to(self.device)
optimizer = optim.Adam(model.parameters(), lr=0.005)
scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=5)
earlystopping = EarlyStopping('min', patience=30)
# check vae dir exists, if not, create it
save_dir = join(args.logdir, 'model')
if not exists(save_dir):
mkdir(save_dir)
reload_file = join(save_dir, 'best.tar')
if not args.noreload and exists(reload_file):
state = torch.load(reload_file)
print("Reloading model at epoch {}"
", with test error {}".format(state['epoch'],
state['precision']))
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
scheduler.load_state_dict(state['scheduler'])
earlystopping.load_state_dict(state['earlystopping'])
# retrain with new dataset with same map dimension
if (args.data_path):
print('TRAIN with new dataset within same geohash range ')
data = pd.read_csv(args.data_path)
data = self.tclass.geohash_decode(data)
data = self.tclass.preprocess(data)
else:
data = self.tclass.saved_data
npdata = data['demand_map'].values
npdata = np.array(list(npdata), dtype=np.float)
tp_data = data[['day', 'timestamp']].values
tp_data = np.array(tp_data, dtype=np.float)
dataset_train = MapDataset(npdata,
tp_data,
self.LOOKBACK,
self.LOOKFORWARD,
train=True)
dataset_test = MapDataset(npdata,
tp_data,
self.LOOKBACK,
self.LOOKFORWARD,
train=False)
train_loader = torch.utils.data.DataLoader(dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
test_loader = torch.utils.data.DataLoader(dataset_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=2)
cur_best = None
best_epoch = None
for epoch in range(1, args.epochs + 1):
train(epoch, model, train_loader, optimizer)
test_loss = test(model, test_loader)
scheduler.step(test_loss)
earlystopping.step(test_loss)
# checkpointing
best_filename = join(save_dir, 'best.tar')
filename = join(save_dir, 'checkpoint.tar')
is_best = not cur_best or test_loss < cur_best
if is_best:
cur_best = test_loss
best_epoch = epoch
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'precision': test_loss,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'earlystopping': earlystopping.state_dict()
}, is_best, filename, best_filename)
if earlystopping.stop:
print("End of Training because of early stopping at epoch {}".
format(epoch))
print(f"Best epoch {best_epoch} with loss :{cur_best}")
break
class VehModel(Model):
def __init__(self, config):
super().__init__(config)
self.dynamics = VehicleTransitionModel().cuda()
self.optim = torch.optim.Adam(self.dynamics.parameters())
self.earlystopping = EarlyStopping(
patience=self._c.early_stop_patience)
self.set_epoch_length()
def preprocess(self, ):
pass
# def get_sample(self):
# return get_dataset_sample(self._dataset)
# danijar style get sample
# yield method
# choose any episode
# why yield episode and not sample.
# while true
# for files in directory:
# if not in cache add to cache
# for i in random set of cache:
# length limitation?
# yield i episode
# while true
# check for files in dir, add new files to cache
# for i in train_steps number of episodes (sampled from episode cache):
# yield a sample of given length
pass
def get_dataset_sample(self, dataset):
sample = next(dataset)
sample = self.preprocess(sample)
return sample
def preprocess(self, e):
bs = self._c.batch_size
e['x'] = torch.Tensor(e['x'].numpy()).reshape(bs, -1).cuda()
e['y'] = torch.Tensor(e['y'].numpy()).reshape(bs, -1).cuda()
e['phases'] = torch.Tensor(e['phases'].numpy()).reshape(bs, -1).cuda()
e['x'][:, [0, 2]] = e['x'][:, [0, 2]] / 200
e['x'][:, [1, 3]] = e['x'][:, [1, 3]] / 35
e['y'][:, 0] = e['y'][:, 0] / 200
e['y'][:, 1] = e['y'][:, 1] / 35
return e
def set_epoch_length(self):
"""
These many number of batches when sampled from the dataset would lead to 1 epoch.
"""
num_episodes = len(self.train_eps)
episode_length = 500
batch_length = self._c.batch_length
batch_size = self._c.batch_size
self.epoch_length = ceil(
num_episodes * (episode_length - (batch_length - 1)) / batch_size)
test_num_episodes = len(self.test_eps)
self.test_epoch_length = ceil(test_num_episodes *
(episode_length -
(batch_length - 1)) / batch_size)
def batch_update_model(self):
# calculate loss
# loss.backward()
# optim.step()
sample = self.get_sample()
loss = self._loss(sample)
loss.backward()
self.optim.step()
def train(self):
cur_best = None
for epoch in range(self._c.epochs):
self.train_dynamics(epoch)
test_loss = self.test()
# scheduler.step(test_loss)
self.earlystopping.step(test_loss)
# checkpointing
best_filename = self._c.logdir / 'best.tar'
filename = self._c.logdir / f'checkpoint_{epoch}.tar'
is_best = not cur_best or test_loss < cur_best
if is_best:
cur_best = test_loss
if is_best or (epoch % 10 == 0):
checkpoint = {
'epoch': epoch,
'state_dict': self.dynamics.state_dict(),
'precision': test_loss,
'optimizer': self.optim.state_dict(),
'earlystopping': self.earlystopping.state_dict(),
# 'scheduler': scheduler.state_dict(),
}
save_checkpoint(checkpoint, is_best, filename, best_filename)
if self.earlystopping.stop:
print("End of Training because of early stopping at epoch {}".
format(epoch))
break
def train_dynamics(self, epoch):
print('=======================> epoch:', epoch)
train_loss = 0
t1 = time.time()
for u in range(self.epoch_length):
s = self.get_dataset_sample(self._dataset)
self.optim.zero_grad()
y_pred = self.dynamics(s['x'], s['phases'])
loss = F.mse_loss(y_pred, s['y'])
loss.backward()
train_loss += loss
self.optim.step()
if (u % int(self.epoch_length / min(self.epoch_length, 20)) == 0):
t2 = time.time()
print(u, round(t2 - t1, 2),
'{:.10f}'.format(loss.item() / self._c.batch_size))
print('====> Epoch: {} Average loss: {:.10f}'.format(
epoch, train_loss / (self.epoch_length * self._c.batch_size)))
def test(self):
self.dynamics.eval()
test_loss = 0
for u in range(self.test_epoch_length):
s = self.get_dataset_sample(self._dataset)
y_pred = self.dynamics(s['x'], s['phases'])
test_loss += F.mse_loss(y_pred, s['y'])
test_loss /= (self.test_epoch_length * self._c.batch_size)
print('====> Test set loss: {:.10f}'.format(test_loss))
print()
return test_loss
def save(self):
raise NotImplementedError
def load(self):
raise NotImplementedError
def _loss(self):
raise NotImplementedError
def create_reconstructions(self):
pass
class DreamerModel(Model):
def __init__(self, config, preprocessing=False):
super().__init__(config)
self.dynamics = ConvTransitionModel2().cuda()
self.optim = torch.optim.Adam(self.dynamics.parameters())
self.earlystopping = EarlyStopping(patience=self._c.early_stop_patience)
self.set_epoch_length()
self.cutoff = None
self.get_dataset_sample = self.get_dataset_sample2 if preprocessing else self.get_dataset_sample1
def preprocess(self,):
pass
# def get_sample(self):
# return get_dataset_sample(self._dataset)
# danijar style get sample
# yield method
# choose any episode
# why yield episode and not sample.
# while true
# for files in directory:
# if not in cache add to cache
# for i in random set of cache:
# length limitation?
# yield i episode
# while true
# check for files in dir, add new files to cache
# for i in train_steps number of episodes (sampled from episode cache):
# yield a sample of given length
pass
def get_dataset_sample_old(self, dataset):
s = next(dataset)
sample = {}
sample['phases'] = self.preprocess(torch.Tensor(s['phases'][:, :, 0, 3, :, 0].numpy())).permute(1, 0, 2)
# sample['y'] = self.preprocess(torch.Tensor(s['x'][:, 1:, 0, :, :, 0].numpy()))
sample['v'] = self.preprocess(torch.Tensor(s['x'][:, :, :, :, :, 1].numpy())) + 0.5
sample['x'] = self.preprocess(torch.Tensor(s['x'][:, :, :, :, :, 0].numpy())).permute(1, 0, 2, 3, 4)
self.cutoff = 0.0
## not needed for now.
# sample['reward'] = s['reward'].numpy()
# sample['action'] = s['action'].numpy()
return sample
def get_dataset_sample1(self, dataset):
# No preprocessing.
s = dataset if isinstance(dataset, dict) else next(dataset)
sample = {}
sample['phases'] = torch.Tensor(s['phases'][:, :, 0, 3, :, 0].numpy()).cuda().permute(1, 0, 2)
sample['x'] = torch.Tensor(s['x'][:, :, :, :, :, 0].numpy()).cuda().permute(1, 0, 2, 3, 4)
sample['action'] = torch.Tensor(s['corrected_action'][:, :].numpy()).cuda().permute(1, 0)
sample['phase_action'] = torch.Tensor(s['corrected_p_action'][:, :].numpy()).cuda().permute(1, 0, 2)
self.cutoff = 0.5
# classification model only works on the last lane
sample['x'] = sample['x'][:, :, 0, -1]
# dreamer requires action indices be 1 step behind observations
sample['x'] = sample['x'][1:]
sample['phases'] = sample['phases'][1:]
sample['action'] = sample['action'][:-1]
sample['phase_action'] = sample['phase_action'][:-1]
# sample['reward'] = s['reward'].numpy()
return sample
def get_dataset_sample2(self, dataset):
# With preprocessing.
s = dataset if isinstance(dataset, dict) else next(dataset)
sample = {}
sample['phases'] = self.preprocess(torch.Tensor(s['phases'][:, :, 0, 3, :, 0].numpy())).permute(1, 0, 2)
sample['x'] = self.preprocess(torch.Tensor(s['x'][:, :, :, :, :, 0].numpy())).permute(1, 0, 2, 3, 4)
sample['action'] = torch.Tensor(s['corrected_action'][:, :].numpy()).cuda().permute(1, 0)
sample['phase_action'] = torch.Tensor(s['corrected_p_action'][:, :].numpy()).cuda().permute(1, 0, 2)
self.cutoff = 0.0
# classification model only works on the last lane
sample['x'] = sample['x'][:, :, 0, -1]
# dreamer requires action indices be 1 step behind observations
sample['x'] = sample['x'][1:]
sample['phases'] = sample['phases'][1:]
sample['action'] = sample['action'][:-1]
sample['phase_action'] = sample['phase_action'][:-1]
# sample['reward'] = s['reward'].numpy()
return sample
def preprocess(self, x):
x = x - 0.5
return x.cuda()
def set_epoch_length(self):
"""
These many number of batches when sampled from the dataset would lead to 1 epoch.
"""
num_episodes = len(self.train_eps)
episode_length = 500
batch_length = self._c.batch_length
batch_size = self._c.batch_size
self.epoch_length = ceil(num_episodes * (episode_length - (batch_length - 1)) / batch_size)
test_num_episodes = len(self.test_eps)
self.test_epoch_length = ceil(test_num_episodes * (episode_length - (batch_length - 1)) / batch_size)
def batch_update_model(self):
# calculate loss
# loss.backward()
# optim.step()
sample = self.get_sample()
loss = self._loss(sample)
loss.backward()
self.optim.step()
def train(self):
cur_best = None
for epoch in range(self._c.epochs):
self.train_dynamics(epoch)
test_loss = self.test()
# scheduler.step(test_loss)
self.earlystopping.step(test_loss)
# checkpointing
best_filename = self._c.logdir / 'best.tar'
filename = self._c.logdir / f'checkpoint_{epoch}.tar'
is_best = not cur_best or test_loss < cur_best
if is_best:
cur_best = test_loss
if is_best or (epoch % 10 == 0):
checkpoint = {
'epoch': epoch,
'state_dict': self.dynamics.state_dict(),
'precision': test_loss,
'optimizer': self.optim.state_dict(),
'earlystopping': self.earlystopping.state_dict(),
# 'scheduler': scheduler.state_dict(),
}
save_checkpoint(checkpoint, is_best, filename, best_filename)
def train_dynamics(self, epoch):
print('=======================> epoch:', epoch)
train_loss = 0
t1 = time.time()
for u in range(self.epoch_length):
s = self.get_dataset_sample(self._dataset)
self.optim.zero_grad()
y_pred = self.dynamics(s['x'], s['phases'])
loss = F.mse_loss(y_pred, s['y'])
loss.backward()
train_loss += loss
self.optim.step()
if (u % int(self.epoch_length/min(self.epoch_length, 20)) == 0):
t2 = time.time()
print(u, round(t2-t1, 2), '{:.10f}'.format(loss.item() / self._c.batch_size))
print('====> Epoch: {} Average loss: {:.10f}'.format(epoch, train_loss / (self.epoch_length * self._c.batch_size)))
def test(self):
self.dynamics.eval()
test_loss = 0
for u in range(self.test_epoch_length):
s = self.get_dataset_sample(self._dataset)
y_pred = self.dynamics(s['x'], s['phases'])
test_loss += F.mse_loss(y_pred, s['y'])
test_loss /= (self.test_epoch_length * self._c.batch_size)
print('====> Test set loss: {:.10f}'.format(test_loss))
print()
return test_loss
def save(self):
raise NotImplementedError
def load(self):
raise NotImplementedError
def _loss(self):
raise NotImplementedError
def create_reconstructions(self):
pass