├── .gitignore ├── README.md ├── main.py ├── models.py ├── readdata.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | data 2 | checkpoints* 3 | *.pyc 4 | __pycache__ 5 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Implementation of [Learning Deep Representation for Trajectory Clustering](https://onlinelibrary.wiley.com/doi/abs/10.1111/exsy.12252) 2 | 3 | This is **NOT** the official implementation. 4 | The code is implemented as one competitor of [Trajectory Similarity Learning with Auxiliary Supervision and Optimal Matching](https://www.ijcai.org/proceedings/2020/0444.pdf). 5 | The problem definition is not exactly same as the original paper, some structure has changed to fit our setting, so this implementation may **NOT** work as expected. 6 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import argparse 3 | import time 4 | import pdb 5 | import os 6 | import pickle 7 | import numpy as np 8 | import time 9 | 10 | import models 11 | import readdata 12 | from utils import cuda, sequence_mask, MSE 13 | 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('-train_data', type=str, default='data/train.pkl') 16 | parser.add_argument('-test_data', type=str, default='data/test.pkl') 17 | parser.add_argument('-eval_size', type=float, default=0.05) 18 | parser.add_argument('-hidden_size', type=int, default=128) 19 | parser.add_argument('-layers', type=int, default=1) 20 | parser.add_argument('-dropout', type=float, default=0.0) 21 | parser.add_argument('-learning_rate', type=float, default=0.001) 22 | parser.add_argument('-batch', type=int, default=256) 23 | parser.add_argument('-epoch', type=int, default=15) 24 | parser.add_argument('-print', type=int, default=20) 25 | parser.add_argument('-save', type=int, default=2000) 26 | parser.add_argument('-max_length', type=int, default=200) 27 | parser.add_argument('-cuda', type=bool, default=True) 28 | parser.add_argument('-checkpoint', type=str, default='checkpoints') 29 | parser.add_argument('-load_state', type=str, default='') 30 | parser.add_argument('-test_when_save', type=bool, default=True) 31 | parser.add_argument('-save_test_result', type=bool, default=True) 32 | parser.add_argument('-encode_savepath', type=str, default='') 33 | 34 | args = parser.parse_args() 35 | 36 | def eval_data(dataset, process = 0): 37 | all_result = [] 38 | all_loss = [] 39 | process = min(process, len(dataset)) 40 | for number, [length, traj, index] in enumerate(dataset): 41 | traj = traj.transpose(0, 1) 42 | fake_input = cuda(torch.zeros((traj.shape[0], traj.shape[1], 0)).float()) 43 | model.eval() 44 | result = model(traj, length, fake_input) 45 | raw_output = model.get_result(traj, length).cpu().detach() 46 | output = torch.tensor(raw_output) 47 | for num in range(len(raw_output)): 48 | output[index[num]] = raw_output[num] 49 | all_result.append(output) 50 | mask = sequence_mask(length, args.max_length).transpose(0, 1) 51 | eval_loss = loss(result, traj, dim = 2) * mask 52 | eval_loss = eval_loss.sum(dim=0) / length.float() 53 | all_loss.append(eval_loss.cpu().detach()) 54 | if process > 0 and number % (len(dataset) // process) == 0: 55 | print('encoding %d / %d' % (number, len(dataset))) 56 | all_result = torch.cat(all_result) 57 | all_loss = torch.cat(all_loss) 58 | return all_result, all_loss.mean().item() 59 | 60 | if args.encode_savepath != '': 61 | start_time = time.time() 62 | cp = torch.load(args.load_state) 63 | c_args = cp['args'] 64 | c_args.load_state = args.load_state 65 | c_args.encode_savepath = args.encode_savepath 66 | c_args.test_data = args.test_data 67 | args = c_args 68 | print('checkpoint arguments loaded') 69 | print(args) 70 | print(time.time() - start_time) 71 | #start_time = time.time() 72 | testdata = readdata.readfile(args.test_data, args.batch, args.max_length, 'cut', False) 73 | model = cuda(models.EncoderDecoder(2, args.hidden_size, args.layers, args.dropout, False), args.cuda) 74 | model.load_state_dict(cp['state_dict']) 75 | loss = MSE 76 | print(time.time() - start_time) 77 | #start_time = time.time() 78 | all_test_result, all_test_loss = eval_data(testdata, 100) 79 | all_test_result = np.array(all_test_result) 80 | pickle.dump(all_test_result, open(args.encode_savepath, 'wb')) 81 | print(time.time() - start_time) 82 | exit() 83 | 84 | print(args) 85 | 86 | if not os.path.exists(args.checkpoint): 87 | os.mkdir(args.checkpoint) 88 | if not os.path.isdir(args.checkpoint): 89 | print('checkpoint path exists and is not a directory! exit.') 90 | exit(0) 91 | 92 | model = cuda(models.EncoderDecoder(2, args.hidden_size, args.layers, args.dropout, False), args.cuda) 93 | traindata, evaldata = readdata.readfile(args.train_data, args.batch, args.max_length, split = 0.05) 94 | testdata = readdata.readfile(args.test_data, args.batch, args.max_length, 'cut', False) 95 | loss = MSE 96 | opt = torch.optim.Adam(model.parameters(), args.learning_rate) 97 | 98 | pdb.set_trace() 99 | 100 | print('train data batch:', len(traindata)) 101 | print('evaluate data batch:', len(evaldata)) 102 | print('test data batch:', len(testdata)) 103 | 104 | epoch = 0 105 | iteration = 0 106 | best_eval_loss = 1e30 107 | best_model_path = os.path.join(args.checkpoint, 'best_model.pt') 108 | if os.path.exists(best_model_path): 109 | best_eval_loss = torch.load(best_model_path)['eval_loss'] 110 | 111 | if args.load_state != '': 112 | checkpoint = torch.load(args.load_state) 113 | epoch = checkpoint['epoch'] - 1 114 | iteration = checkpoint['iteration'] 115 | checkpoint_args = checkpoint['args'] 116 | model.load_state_dict(checkpoint['state_dict']) 117 | opt.load_state_dict(checkpoint['opt_state_dict']) 118 | 119 | ite_start_time = time.time() 120 | for epoch in range(epoch + 1, args.epoch + 1): 121 | print('start epoch %4d' % epoch) 122 | for length, traj, index in traindata: 123 | traj = traj.transpose(0, 1) 124 | fake_input = cuda(torch.zeros((traj.shape[0], traj.shape[1], 0)).float()) 125 | model.train() 126 | opt.zero_grad() 127 | result = model(traj, length, fake_input) 128 | mask = sequence_mask(length, args.max_length).transpose(0, 1) 129 | l = loss(result, traj, dim = 2) * mask 130 | l = (l.sum(dim=0) / length.float()).mean() 131 | l.backward() 132 | opt.step() 133 | iteration += 1 134 | if iteration % args.print == 0: 135 | model.eval() 136 | _, eval_loss = eval_data(evaldata) 137 | print('iteration %d: train_loss:%.10f eval_loss:%.10f time:%.2f' % (iteration, l, eval_loss, time.time() - ite_start_time)) 138 | ite_start_time = time.time() 139 | if iteration % args.save == 0: 140 | _, all_eval_loss = eval_data(evaldata) 141 | all_test_result = None 142 | all_test_loss = None 143 | if args.test_when_save: 144 | all_test_result, all_test_loss = eval_data(testdata) 145 | if not args.save_test_result: 146 | all_test_result = None 147 | filename = '%06d.pt' % iteration 148 | filename = os.path.join(args.checkpoint, filename) 149 | savedata = { 150 | 'epoch': epoch, 151 | 'iteration': iteration, 152 | 'train_loss': l, 153 | 'eval_loss': all_eval_loss, 154 | 'test_loss': all_test_loss, 155 | 'test_result': all_test_result, 156 | 'args': args, 157 | 'state_dict': model.state_dict(), 158 | 'opt_state_dict': opt.state_dict() 159 | } 160 | torch.save(savedata, filename) 161 | print('checkpoint %s saved. test_loss: %s' % (filename, str(all_test_loss))) 162 | if all_eval_loss < best_eval_loss: 163 | best_eval_loss = all_eval_loss 164 | torch.save(savedata, best_model_path) 165 | print('best model updated') 166 | 167 | -------------------------------------------------------------------------------- /models.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | from torch.nn.utils.rnn import pad_packed_sequence 5 | from torch.nn.utils.rnn import pack_padded_sequence 6 | import os 7 | 8 | class StackingGRUCell(nn.Module): 9 | """ 10 | Multi-layer CRU Cell 11 | """ 12 | def __init__(self, input_size, hidden_size, num_layers, dropout): 13 | super(StackingGRUCell, self).__init__() 14 | self.num_layers = num_layers 15 | self.grus = nn.ModuleList() 16 | self.dropout = nn.Dropout(dropout) 17 | 18 | self.grus.append(nn.GRUCell(input_size, hidden_size)) 19 | for i in range(1, num_layers): 20 | self.grus.append(nn.GRUCell(hidden_size, hidden_size)) 21 | 22 | def forward(self, input, h0): 23 | """ 24 | Input: 25 | input (batch, input_size): input tensor 26 | h0 (num_layers, batch, hidden_size): initial hidden state 27 | --- 28 | Output: 29 | output (batch, hidden_size): the final layer output tensor 30 | hn (num_layers, batch, hidden_size): the hidden state of each layer 31 | """ 32 | hn = [] 33 | output = input 34 | for i, gru in enumerate(self.grus): 35 | hn_i = gru(output, h0[i]) 36 | hn.append(hn_i) 37 | if i != self.num_layers - 1: 38 | output = self.dropout(hn_i) 39 | else: 40 | output = hn_i 41 | hn = torch.stack(hn) 42 | return output, hn 43 | 44 | class GlobalAttention(nn.Module): 45 | """ 46 | $$a = \sigma((W_1 q)H)$$ 47 | $$c = \tanh(W_2 [a H, q])$$ 48 | """ 49 | def __init__(self, hidden_size): 50 | super(GlobalAttention, self).__init__() 51 | self.L1 = nn.Linear(hidden_size, hidden_size, bias=False) 52 | self.L2 = nn.Linear(2*hidden_size, hidden_size, bias=False) 53 | self.softmax = nn.Softmax(dim=1) 54 | self.tanh = nn.Tanh() 55 | 56 | def forward(self, q, H): 57 | """ 58 | Input: 59 | q (batch, hidden_size): query 60 | H (batch, seq_len, hidden_size): context 61 | --- 62 | Output: 63 | c (batch, hidden_size) 64 | """ 65 | # (batch, hidden_size) => (batch, hidden_size, 1) 66 | q1 = self.L1(q).unsqueeze(2) 67 | # (batch, seq_len) 68 | a = torch.bmm(H, q1).squeeze(2) 69 | a = self.softmax(a) 70 | # (batch, seq_len) => (batch, 1, seq_len) 71 | a = a.unsqueeze(1) 72 | # (batch, hidden_size) 73 | c = torch.bmm(a, H).squeeze(1) 74 | # (batch, hidden_size * 2) 75 | c = torch.cat([c, q], 1) 76 | return self.tanh(self.L2(c)) 77 | 78 | class Encoder(nn.Module): 79 | def __init__(self, input_size, hidden_size, num_layers, dropout, 80 | bidirectional): 81 | """ 82 | embedding (vocab_size, input_size): pretrained embedding 83 | """ 84 | super(Encoder, self).__init__() 85 | self.num_directions = 2 if bidirectional else 1 86 | assert hidden_size % self.num_directions == 0 87 | self.hidden_size = hidden_size // self.num_directions 88 | self.num_layers = num_layers 89 | 90 | self.rnn = nn.GRU(input_size, self.hidden_size, 91 | num_layers=num_layers, 92 | bidirectional=bidirectional, 93 | dropout=dropout) 94 | 95 | def forward(self, input, lengths, h0=None): 96 | """ 97 | Input: 98 | input (seq_len, batch): padded sequence tensor 99 | lengths (1, batch): sequence lengths 100 | h0 (num_layers*num_directions, batch, hidden_size): initial hidden state 101 | --- 102 | Output: 103 | hn (num_layers*num_directions, batch, hidden_size): 104 | the hidden state of each layer 105 | output (seq_len, batch, hidden_size*num_directions): output tensor 106 | """ 107 | # (seq_len, batch) => (seq_len, batch, input_size) 108 | embed = input 109 | lengths = lengths.data.view(-1).tolist() 110 | if lengths is not None: 111 | embed = pack_padded_sequence(embed, lengths) 112 | output, hn = self.rnn(embed, h0) 113 | if lengths is not None: 114 | output = pad_packed_sequence(output)[0] 115 | return hn, output 116 | 117 | class Decoder(nn.Module): 118 | def __init__(self, input_size, hidden_size, num_layers, dropout): 119 | super(Decoder, self).__init__() 120 | self.rnn = StackingGRUCell(input_size, hidden_size, num_layers, 121 | dropout) 122 | #self.attention = GlobalAttention(hidden_size) 123 | self.dropout = nn.Dropout(dropout) 124 | self.num_layers = num_layers 125 | 126 | def forward(self, input, h, H, use_attention=False): 127 | """ 128 | Input: 129 | input (seq_len, batch): padded sequence tensor 130 | h (num_layers, batch, hidden_size): input hidden state 131 | H (seq_len, batch, hidden_size): the context used in attention mechanism 132 | which is the output of encoder 133 | use_attention: If True then we use attention 134 | --- 135 | Output: 136 | output (seq_len, batch, hidden_size) 137 | h (num_layers, batch, hidden_size): output hidden state, 138 | h may serve as input hidden state for the next iteration, 139 | especially when we feed the word one by one (i.e., seq_len=1) 140 | such as in translation 141 | """ 142 | embed = input 143 | output = [] 144 | # split along the sequence length dimension 145 | for e in embed.split(1): 146 | e = e.squeeze(0) # (1, batch, input_size) => (batch, input_size) 147 | o, h = self.rnn(e, h) 148 | #if use_attention: 149 | # o = self.attention(o, H.transpose(0, 1)) 150 | o = self.dropout(o) 151 | output.append(o) 152 | output = torch.stack(output) 153 | return output, h 154 | 155 | class EncoderDecoder(nn.Module): 156 | def __init__(self, input_size, hidden_size, num_layers, dropout, bidirectional): 157 | super(EncoderDecoder, self).__init__() 158 | self.input_size = input_size 159 | self.encoder = Encoder(input_size, hidden_size, num_layers, 160 | dropout, bidirectional) 161 | self.decoder = Decoder(0, hidden_size, num_layers, 162 | dropout) 163 | self.linear = nn.Linear(hidden_size, input_size) 164 | self.num_layers = num_layers 165 | 166 | def encoder_hn2decoder_h0(self, h): 167 | """ 168 | Input: 169 | h (num_layers * num_directions, batch, hidden_size): encoder output hn 170 | --- 171 | Output: 172 | h (num_layers, batch, hidden_size * num_directions): decoder input h0 173 | """ 174 | if self.encoder.num_directions == 2: 175 | num_layers, batch, hidden_size = h.size(0)//2, h.size(1), h.size(2) 176 | return h.view(num_layers, 2, batch, hidden_size)\ 177 | .transpose(1, 2).contiguous()\ 178 | .view(num_layers, batch, hidden_size * 2) 179 | else: 180 | return h 181 | 182 | def forward(self, src, lengths, trg): 183 | """ 184 | Input: 185 | src (src_seq_len, batch): source tensor 186 | lengths (1, batch): source sequence lengths 187 | trg (trg_seq_len, batch): target tensor, the `seq_len` in trg is not 188 | necessarily the same as that in src 189 | --- 190 | Output: 191 | output (trg_seq_len, batch, hidden_size) 192 | """ 193 | encoder_hn, H = self.encoder(src, lengths) 194 | decoder_h0 = self.encoder_hn2decoder_h0(encoder_hn) 195 | ## for target we feed the range [BOS:EOS-1] into decoder 196 | output, decoder_hn = self.decoder(trg, decoder_h0, H) 197 | output = self.linear(output) 198 | return output 199 | 200 | def get_result(self, src, lengths): 201 | encoder_hn, H = self.encoder(src, lengths) 202 | decoder_h0 = self.encoder_hn2decoder_h0(encoder_hn).squeeze(0) 203 | return decoder_h0 204 | -------------------------------------------------------------------------------- /readdata.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import random 4 | from torch.utils.data import DataLoader, Dataset 5 | from utils import cuda 6 | import pickle 7 | 8 | class TrajDataset(Dataset): 9 | def __init__(self, trajs, remove_or_cut, max_length = 10000000): 10 | super(TrajDataset, self).__init__() 11 | assert(remove_or_cut == 'remove' or remove_or_cut == 'cut') 12 | self.rawdata = trajs 13 | select = [] 14 | for traj in trajs: 15 | if len(traj) <= max_length: 16 | select.append(traj) 17 | elif remove_or_cut == 'cut': 18 | select.append(traj[:max_length]) 19 | trajs = select 20 | self.length = torch.tensor(list(map(len, trajs))) 21 | self.min = [] 22 | self.max = [] 23 | for i in range(len(trajs[0][0])): 24 | onedim = [[y[i] for y in x] for x in trajs] 25 | self.min.append(min(map(min, onedim))) 26 | self.max.append(max(map(max, onedim))) 27 | self.min = torch.tensor(self.min) 28 | self.max = torch.tensor(self.max) 29 | #print(self.min[0].item(), self.min[1].item(), self.max[0].item(), self.max[1].item()) 30 | #shanghai 31 | #self.min = torch.tensor([30.90562629699707, 121.08155822753906]) 32 | #self.max = torch.tensor([31.45410919189453, 121.81672668457031]) 33 | #porto 34 | self.min = torch.tensor([41.1006965637207, -8.699525833129883]) 35 | self.max = torch.tensor([41.19973373413086, -8.498448371887207]) 36 | self.trajs = self.min.unsqueeze(0).unsqueeze(0).repeat(len(trajs), max_length, 1) 37 | for i in range(len(trajs)): 38 | self.trajs[i][:self.length[i]] = torch.tensor(trajs[i]) 39 | self.trajs -= self.min 40 | self.trajs /= self.max - self.min 41 | def __len__(self): 42 | return len(self.trajs) 43 | def __getitem__(self, x): 44 | return self.length[x], self.trajs[x] 45 | def collate_fn(self, data): 46 | for num, one in enumerate(data): 47 | data[num] = list(one) 48 | data[num].append(num) 49 | data.sort(key=lambda x:-x[0]) 50 | x, y, index = zip(*data) 51 | return cuda(torch.stack(x)), cuda(torch.stack(y)), index 52 | 53 | def readfile(filename, batch_size, max_length, remove_or_cut = 'remove', shuffle = True, split = 0.0): 54 | arr = pickle.load(open(filename, 'rb')) 55 | if split != 0.0: 56 | arr1 = [] 57 | arr2 = [] 58 | index = list(range(len(arr))) 59 | if shuffle: 60 | random.shuffle(index) 61 | for i in range(len(arr)): 62 | if index[i] < split * len(arr): 63 | arr2.append(arr[i]) 64 | else: 65 | arr1.append(arr[i]) 66 | dataset1 = TrajDataset(arr1, remove_or_cut, max_length) 67 | dataset2 = TrajDataset(arr2, remove_or_cut, max_length) 68 | return DataLoader(dataset1, batch_size, shuffle, collate_fn=dataset1.collate_fn), DataLoader(dataset2, batch_size, shuffle, collate_fn=dataset2.collate_fn) 69 | dataset = TrajDataset(arr, remove_or_cut, max_length) 70 | return DataLoader(dataset, batch_size, shuffle, collate_fn=dataset.collate_fn) 71 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | last_use_cuda = True 4 | 5 | def cuda(tensor, use_cuda = None): 6 | """ 7 | A cuda wrapper 8 | """ 9 | global last_use_cuda 10 | if use_cuda == None: 11 | use_cuda = last_use_cuda 12 | last_use_cuda = use_cuda 13 | if not use_cuda: 14 | return tensor 15 | if tensor is None: 16 | return None 17 | if torch.cuda.is_available(): 18 | return tensor.cuda() 19 | else: 20 | return tensor 21 | 22 | def sequence_mask(sequence_length, max_length=None): 23 | """ 24 | e.g., sequence_length = "5,7,8", max_length=None 25 | it will return 26 | tensor([[ 1, 1, 1, 1, 1, 0, 0, 0], 27 | [ 1, 1, 1, 1, 1, 1, 1, 0], 28 | [ 1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.float32) 29 | :param sequence_length: a torch tensor 30 | :param max_length: if not given, it will be set to the maximum of `sequence_length` 31 | :return: a tensor with dimension [*sequence_length.size(), max_length] 32 | """ 33 | if len(sequence_length.size()) > 1: 34 | ori_shape = list(sequence_length.size()) 35 | sequence_length = sequence_length.view(-1) # [N, ?] 36 | reshape_back = True 37 | else: 38 | reshape_back = False 39 | 40 | if max_length is None: 41 | max_length = sequence_length.data.max() 42 | batch_size = sequence_length.size(0) 43 | seq_range = torch.arange(0, max_length).long() # [max_length] 44 | seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_length) # [batch, max_len], repeats on each column 45 | seq_range_expand = torch.autograd.Variable(seq_range_expand).to(sequence_length.device) 46 | #if sequence_length.is_cuda: 47 | # seq_range_expand = seq_range_expand.cuda() 48 | seq_length_expand = sequence_length.unsqueeze(1).expand_as(seq_range_expand) # [batch, max_len], repeats on each row 49 | 50 | ret = (seq_range_expand < seq_length_expand).float() # [batch, max_len] 51 | 52 | if reshape_back: 53 | ret = ret.view(ori_shape + [max_length]) 54 | 55 | return ret 56 | 57 | def MSE(src, dest, dim = None): 58 | res = src - dest 59 | res = res * res 60 | if dim == None: 61 | return res.mean() 62 | return res.mean(dim = dim) 63 | --------------------------------------------------------------------------------