Example #1
0
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
Example #2
0
    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
Example #3
0
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
Example #4
0
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