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