class DataSource(object): def __init__(self, mode, data_path): self.normalizer = Normalizer() self.data_path = data_path self.mode = mode self.filenames = [] self.contents = [] self.labels = [] def load_data(self): # load raw sample raw_data = [] raw_sample = "" with open(self.data_path, "r") as f: for row in f: if (self.mode + "_") in row: raw_data.append(raw_sample) raw_sample = row else: raw_sample += row raw_data.append(raw_sample) raw_data = raw_data[1:] # process raw data for i, raw_sample in enumerate(raw_data): self.filenames.append(raw_sample.split("\n")[0]) print(i, end = "\r") content = re.search("\"(.*)\"", raw_sample, re.DOTALL).group(1) self.contents.append(self.normalizer.transform(content)) if self.mode == "train": self.labels.append(raw_sample.split("\n")[-3])
class Trainer: loss_dispatcher = {"l2": nn.MSELoss} def __init__(self, **params): self.num_epochs = params["num_epochs"] self.early_stop_tolerance = params["early_stop_tolerance"] self.norm_method = params["norm_method"] self.loss_type = params["loss_type"] self.learning_rate = params["learning_rate"] self.l2_reg = params["l2_reg"] self.clip = params['clip'] self.device = params['device'] self.input_normalizer = Normalizer(self.norm_method) self.output_normalizer = Normalizer(self.norm_method) def fit(self, model, batch_generator): """ :param TrajGRU model: :param BatchGenerator batch_generator: :return: """ print('Training starts...') model.to(self.device) train_loss = [] val_loss = [] tolerance = 0 best_epoch = 0 best_val_loss = 1e6 evaluation_val_loss = best_val_loss best_dict = model.state_dict() data_list = [] label_list = [] for x, y in batch_generator.generate('train'): data_list.append(x.reshape(-1, *x.shape[2:])) label_list.append(y.reshape(-1, *y.shape[2:])) self.input_normalizer.fit(torch.cat(data_list)) self.output_normalizer.fit(torch.cat(label_list)) optimizer = optim.Adam(model.parameters(), lr=self.learning_rate, weight_decay=self.l2_reg) for epoch in range(self.num_epochs): # train and validation loop start_time = time.time() # train running_train_loss = self.step_loop( model=model, batch_generator=batch_generator, step_fun=self.train_step, loss_fun=self.loss_fun, dataset_type='train', optimizer=optimizer, denormalize=False) # validate running_val_loss = self.step_loop(model=model, batch_generator=batch_generator, step_fun=self.eval_step, loss_fun=self.loss_fun, dataset_type='validation', optimizer=None, denormalize=False) epoch_time = time.time() - start_time message_str = "\nEpoch: {}, Train_loss: {:.5f}, Validation_loss: {:.5f}, Took {:.3f} seconds." print( message_str.format(epoch + 1, running_train_loss, running_val_loss, epoch_time)) # save the losses train_loss.append(running_train_loss) val_loss.append(running_val_loss) if running_val_loss < best_val_loss: best_epoch = epoch + 1 best_val_loss = running_val_loss best_dict = deepcopy(model.state_dict()) # brutal tolerance = 0 else: tolerance += 1 if tolerance > self.early_stop_tolerance or epoch == self.num_epochs - 1: model.load_state_dict(best_dict) evaluation_val_loss = self.step_loop( model=model, batch_generator=batch_generator, step_fun=self.eval_step, loss_fun=self.loss_fun_evaluation, dataset_type='validation', optimizer=None, denormalize=True) message_str = "Early exiting from epoch: {}, Validation error: {:.5f}." print(message_str.format(best_epoch, evaluation_val_loss)) break print('Training finished') return train_loss, val_loss, evaluation_val_loss def step_loop(self, model, batch_generator, step_fun, loss_fun, dataset_type, optimizer, denormalize): count = 0 running_loss = 0.0 dataset = batch_generator.dataset_dict[dataset_type] hidden = self.reset_per_epoch(model=model, batch_size=batch_generator.batch_size) for count, (input_data, output_data) in enumerate( batch_generator.generate(dataset_type)): print("\r{:.2f}%".format(dataset.count * 100 / len(dataset)), flush=True, end='') input_data_shape = input_data.shape input_data = self.input_normalizer.transform( input_data.reshape(-1, *input_data.shape[2:])) input_data = input_data.reshape(input_data_shape).to(self.device) output_data_shape = output_data.shape output_data = self.output_normalizer.transform( output_data.reshape(-1, *output_data.shape[2:])) output_data = output_data.reshape(output_data_shape).to( self.device) loss = step_fun(model=model, input_tensor=input_data, output_tensor=output_data, hidden=hidden, loss_fun=loss_fun, optimizer=optimizer, denormalize=denormalize) # many-to-one hidden = self.repackage_hidden(hidden) running_loss += loss.item() running_loss /= (count + 1) return running_loss def train_step(self, model, input_tensor, output_tensor, hidden, loss_fun, optimizer, denormalize): def closure(): optimizer.zero_grad() predictions = model.forward(input_tensor, hidden) loss = loss_fun(predictions, output_tensor) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), self.clip) return loss loss = optimizer.step(closure) return loss def eval_step(self, model, input_tensor, output_tensor, hidden, loss_fun, optimizer, denormalize): predictions = model.forward(input_tensor, hidden) if denormalize: predictions = self.output_normalizer.inverse_transform( predictions.to('cpu')) output_tensor = self.output_normalizer.inverse_transform( output_tensor.to('cpu')) # output_tensor = output_tensor.to('cpu') loss = loss_fun(predictions, output_tensor) return loss def loss_fun(self, predictions, labels): """ :param predictions: BxD_out :param labels: BxD_out :return: """ loss_obj = self.loss_dispatcher[self.loss_type]() loss = loss_obj(predictions, labels) return loss def loss_fun_evaluation(self, predictions, labels): """ :param labels: :param preds: :return: """ predictions = predictions.detach().numpy() labels = labels.detach().numpy() loss = haversine_dist(predictions, labels).mean() return loss def reset_per_epoch(self, model, batch_size): """ This will be called at beginning of every epoch :param model: :param batch_size: :return: """ hidden_list = model.init_hidden(batch_size=batch_size, device=self.device) return hidden_list def repackage_hidden(self, h): """ Wraps hidden states in new Tensors, to detach them from their history. :param h: list of states, e.g [state, state, ...] """ if isinstance(h, torch.Tensor): return h.detach() else: return tuple(self.repackage_hidden(v) for v in h)
class LSTM(nn.Module): optimizer_dispatcher = {"adam": torch.optim.Adam, "sgd": torch.optim.SGD} loss_dispatcher = {"l2": nn.MSELoss} activation_dispatcher = { "leaky_relu": nn.LeakyReLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid } def __init__(self, input_dim, output_dim, **params): super(LSTM, self).__init__() self.input_dim = input_dim self.output_dim = output_dim self.loss_type = params["loss_type"] self.learning_rate = params["learning_rate"] self.optimizer_type = params["optimizer_type"] self.grad_clip = params["grad_clip"] self.l2_reg = params["l2_reg"] self.dropout_rate = params["dropout_rate"] self.num_epochs = params["num_epochs"] self.early_stop_tolerance = params["early_stop_tolerance"] self.hidden_dim_list = params["hidden_dim_list"] self.num_layers = len(self.hidden_dim_list) self.final_act_type = params["final_act_type"] self.relu_alpha = params["relu_alpha"] self.input_norm_method = params["input_norm_method"] self.output_norm_method = params["output_norm_method"] self.__create_rnn_cell_list() self.__create_dropout_layer() self.__create_dense_layer() self.__create_final_act_layer() self.optimizer = self.optimizer_dispatcher[self.optimizer_type]( self.parameters(), lr=self.learning_rate, weight_decay=self.l2_reg) self.input_normalizer = Normalizer(self.input_norm_method) self.output_normalizer = Normalizer(self.output_norm_method) def __create_rnn_cell_list(self): cell_list = [] input_dim = self.input_dim for i in range(self.num_layers): hidden_dim = self.hidden_dim_list[i] cell_list.append( nn.LSTMCell(input_size=input_dim, hidden_size=hidden_dim)) input_dim = self.hidden_dim_list[i] self.cell_list = nn.ModuleList(cell_list) def __create_dropout_layer(self): self.dropout_layer = nn.Dropout(p=self.dropout_rate) def __create_dense_layer(self): self.dense_layer = nn.Linear(in_features=self.hidden_dim_list[-1], out_features=self.output_dim) def __create_final_act_layer(self): if self.final_act_type == "leaky_relu": self.final_act_layer = self.activation_dispatcher[ self.final_act_type](self.relu_alpha) else: self.final_act_layer = self.activation_dispatcher[ self.final_act_type]() def __init_hidden_states(self, batch_size): self.h_list = [] self.c_list = [] for cell in self.cell_list: self.h_list.append( Variable(torch.zeros(batch_size, cell.hidden_size))) self.c_list.append( Variable(torch.zeros(batch_size, cell.hidden_size))) @staticmethod def repackage_hidden(h): """Wraps hidden states in new Tensors, to detach them from their history.""" if isinstance(h, torch.Tensor): return h.detach() else: return tuple(LSTM.repackage_hidden(v) for v in h) def forward(self, input_tensor): batch_size = input_tensor.shape[0] num_steps = input_tensor.shape[1] self.__init_hidden_states(batch_size) for step in range(num_steps): x = input_tensor[:, step] for layer_idx, cell in enumerate(self.cell_list): h, c = cell(x, (self.h_list[layer_idx], self.c_list[layer_idx])) self.h_list[layer_idx] = h self.c_list[layer_idx] = c x = h x = self.dropout_layer(x) x = self.dense_layer(x) x = self.final_act_layer(x) return x def fit(self, batch_generator): print('Training starts...') train_loss = [] val_loss = [] tolerance = 0 best_epoch = 0 best_val_loss = 1e6 evaluation_val_loss = best_val_loss best_dict = self.state_dict() self.input_normalizer.fit(batch_generator.dataset_dict["train"].data) self.output_normalizer.fit(batch_generator.dataset_dict["train"].label) for epoch in range(self.num_epochs): # train and validation loop start_time = time.time() running_train_loss = self.step_loop(batch_generator, self.train_step, self.loss_fun, 'train', denormalize=False) running_val_loss = self.step_loop(batch_generator, self.eval_step, self.loss_fun, 'validation', denormalize=False) epoch_time = time.time() - start_time message_str = "Epoch: {}, Train_loss: {:.5f}, Validation_loss: {:.5f}, Took {:.3f} seconds." print( message_str.format(epoch + 1, running_train_loss, running_val_loss, epoch_time)) # save the losses train_loss.append(running_train_loss) val_loss.append(running_val_loss) if running_val_loss < best_val_loss: best_epoch = epoch + 1 best_val_loss = running_val_loss best_dict = deepcopy(self.state_dict()) # brutal tolerance = 0 else: tolerance += 1 if tolerance > self.early_stop_tolerance or epoch == self.num_epochs - 1: self.load_state_dict(best_dict) evaluation_val_loss = self.step_loop(batch_generator, self.eval_step, self.loss_fun_evaluation, 'validation', denormalize=True) message_str = "Early exiting from epoch: {}, Validation error: {:.5f}." print(message_str.format(best_epoch, evaluation_val_loss)) break print('Training finished') return train_loss, val_loss, evaluation_val_loss def step_loop(self, batch_generator, step_fun, loss_fun, dataset_type, denormalize): count = 0 running_loss = 0.0 for count, (input_data, output_data) in enumerate( batch_generator.generate(dataset_type)): input_data = self.input_normalizer.transform(input_data) output_data = self.output_normalizer.transform(output_data) loss = step_fun(input_data, output_data[:, -1], loss_fun, denormalize) # many-to-one try: running_loss += loss.detach().numpy() except: running_loss += loss running_loss /= (count + 1) return running_loss def train_step(self, input_tensor, output_tensor, loss_fun, denormalize): def closure(): self.optimizer.zero_grad() predictions = self.forward(input_tensor) loss = loss_fun(predictions, output_tensor) loss.backward() nn.utils.clip_grad_norm_(self.parameters(), self.grad_clip) return loss loss = self.optimizer.step(closure) return loss def eval_step(self, input_tensor, output_tensor, loss_fun, denormalize): predictions = self.forward(input_tensor) if denormalize: predictions = self.output_normalizer.inverse_transform(predictions) output_tensor = self.output_normalizer.inverse_transform( output_tensor) loss = loss_fun(predictions, output_tensor) return loss def loss_fun(self, predictions, labels): """ :param predictions: BxD_out :param labels: BxD_out :return: """ loss_obj = self.loss_dispatcher[self.loss_type]() loss = loss_obj(predictions, labels) return loss def loss_fun_evaluation(self, predictions, labels): """ :param labels: :param preds: :return: """ predictions = predictions.detach().numpy() labels = labels.detach().numpy() loss = haversine_dist(predictions, labels).mean() return loss