class TensorboardLogger(BaseLogger): def open(self, root_dir: str): from pathlib import Path log_dir = Path(root_dir, 'tb_logs') self.logger = SummaryWriter(str(log_dir)) def close(self): self.logger.close() def log_scalar(self, name: str, value: float, step: Optional[int] = None): self.logger.add_scalar(name, value, step) def log_scalars(self, name: str, scalar_dict: Dict[str, float], step: Optional[int] = None): self.logger.add_scalars(name, scalar_dict, step) def log_model(self, model: ITrainableModel, device: str): inputs = model.example_inputs() if inputs is None: pass elif isinstance(inputs, torch.Tensor): inputs = inputs.to(device) elif isinstance(inputs, (list, tuple)): inputs = [x.to(device) for x in inputs] elif isinstance(inputs, dict): inputs = {k: v.to(device) for k, v in inputs.items()} else: raise RuntimeError('Unsupported example_inputs') self.logger.add_graph(model, input_to_model=inputs) def log_info(self, info: str): pass
class SeqClassificationModel(Model): """ Question answering model where answers are sentences """ def __init__( self, vocab: Vocabulary, text_field_embedder: TextFieldEmbedder, use_sep: bool = True, with_crf: bool = False, self_attn: Seq2SeqEncoder = None, bert_dropout: float = 0.1, sci_sum: bool = False, additional_feature_size: int = 0, ) -> None: super(SeqClassificationModel, self).__init__(vocab) self.track_embedding_list = [] self.track_embedding = {} self.text_field_embedder = text_field_embedder self.vocab = vocab self.use_sep = use_sep self.with_crf = with_crf self.sci_sum = sci_sum self.self_attn = self_attn self.additional_feature_size = additional_feature_size self.dropout = torch.nn.Dropout(p=bert_dropout) # define loss if self.sci_sum: self.loss = torch.nn.MSELoss( reduction='none') # labels are rouge scores self.labels_are_scores = True self.num_labels = 1 else: self.loss = torch.nn.CrossEntropyLoss(ignore_index=-1, reduction='none') self.labels_are_scores = False self.num_labels = self.vocab.get_vocab_size(namespace='labels') # define accuracy metrics self.label_accuracy = CategoricalAccuracy() self.label_f1_metrics = {} # define F1 metrics per label for label_index in range(self.num_labels): label_name = self.vocab.get_token_from_index( namespace='labels', index=label_index) self.label_f1_metrics[label_name] = F1Measure(label_index) encoded_senetence_dim = text_field_embedder._token_embedders[ 'bert'].output_dim ff_in_dim = encoded_senetence_dim if self.use_sep else self_attn.get_output_dim( ) ff_in_dim += self.additional_feature_size self.time_distributed_aggregate_feedforward = TimeDistributed( Linear(ff_in_dim, self.num_labels)) if self.with_crf: self.crf = ConditionalRandomField( self.num_labels, constraints=None, include_start_end_transitions=True) self.track_embedding["init_info"] = { "ff_in_dim": ff_in_dim, "encoded_sentence_dim": encoded_senetence_dim, "sci_sum": self.sci_sum, "use_sep": self.use_sep, "with_crf": self.with_crf, "additional_feature_size": self.additional_feature_size } self.t_board_writer = SummaryWriter() self.t_board_writer.add_graph(self) def forward( self, # type: ignore sentences: torch.LongTensor, labels: torch.IntTensor = None, confidences: torch.Tensor = None, additional_features: torch.Tensor = None, ) -> Dict[str, torch.Tensor]: # pylint: disable=arguments-differ """ Parameters ---------- TODO: add description Returns ------- An output dictionary consisting of: loss : torch.FloatTensor, optional A scalar loss to be optimised. """ # =========================================================================================================== # Layer 1: For each sentence, participant pair: create a Glove embedding for each token # Input: sentences # Output: embedded_sentences print(sentences) sentences_conv = {} for key, val in sentences_conv.items(): sentences_conv[key] = val.cpu().data.numpy().tolist() self.track_embedding["Transformation_0"] = { "sentences": sentences_conv } # embedded_sentences: batch_size, num_sentences, sentence_length, embedding_size embedded_sentences = self.text_field_embedder(sentences) self.track_embedding["Transformation_1"] = { "size": list(embedded_sentences.size()), "dim": embedded_sentences.dim() } # Kacper: Basically a padding mask for bert mask = get_text_field_mask(sentences, num_wrapping_dims=1).float() batch_size, num_sentences, _, _ = list(embedded_sentences.size()) if self.use_sep: # The following code collects vectors of the SEP tokens from all the examples in the batch, # and arrange them in one list. It does the same for the labels and confidences. # TODO: replace 103 with '[SEP]' # Kacper: This is an important step where we get SEP tokens to later do sentence classification # Kacper: We take a location of SEP tokens from the sentences to get a mask sentences_mask = sentences[ 'bert'] == 103 # mask for all the SEP tokens in the batch # Kacper: We use this mask to get the respective embeddings from the output layer of bert embedded_sentences = embedded_sentences[ sentences_mask] # given batch_size x num_sentences_per_example x sent_len x vector_len # returns num_sentences_per_batch x vector_len self.track_embedding["Transformation_2"] = { "size": list(embedded_sentences.size()), "dim": embedded_sentences.dim() } # Kacper: I dont get it why it became 2 instead of 4? What is the difference between size() and dim()??? assert embedded_sentences.dim() == 2 num_sentences = embedded_sentences.shape[0] # Kacper: comment below is vague # Kacper: I think we batch in one array because we just need to compute a mean loss from all of them # for the rest of the code in this model to work, think of the data we have as one example # with so many sentences and a batch of size 1 batch_size = 1 embedded_sentences = embedded_sentences.unsqueeze( dim=0) # Kacper: We batch all sentences in one array self.track_embedding["Transformation_3"] = { "size": list(embedded_sentences.size()), "dim": embedded_sentences.dim() } # Kacper: Dropout layer is between filtered embeddings and linear layer embedded_sentences = self.dropout(embedded_sentences) self.track_embedding["Transformation_4"] = { "size": list(embedded_sentences.size()), "dim": embedded_sentences.dim() } # Kacper: we provide the labels for training (for each sentence) if labels is not None: if self.labels_are_scores: labels_mask = labels != 0.0 # mask for all the labels in the batch (no padding) else: labels_mask = labels != -1 # mask for all the labels in the batch (no padding) labels = labels[ labels_mask] # given batch_size x num_sentences_per_example return num_sentences_per_batch assert labels.dim() == 1 if confidences is not None: confidences = confidences[labels_mask] assert confidences.dim() == 1 if additional_features is not None: additional_features = additional_features[labels_mask] assert additional_features.dim() == 2 num_labels = labels.shape[0] # Kacper: this might be useful to consider in my code as well if num_labels != num_sentences: # bert truncates long sentences, so some of the SEP tokens might be gone assert num_labels > num_sentences # but `num_labels` should be at least greater than `num_sentences` logger.warning( f'Found {num_labels} labels but {num_sentences} sentences' ) labels = labels[: num_sentences] # Ignore some labels. This is ok for training but bad for testing. # We are ignoring this problem for now. # TODO: fix, at least for testing # do the same for `confidences` if confidences is not None: num_confidences = confidences.shape[0] if num_confidences != num_sentences: assert num_confidences > num_sentences confidences = confidences[:num_sentences] # and for `additional_features` if additional_features is not None: num_additional_features = additional_features.shape[0] if num_additional_features != num_sentences: assert num_additional_features > num_sentences additional_features = additional_features[: num_sentences] # similar to `embedded_sentences`, add an additional dimension that corresponds to batch_size=1 labels = labels.unsqueeze(dim=0) if confidences is not None: confidences = confidences.unsqueeze(dim=0) if additional_features is not None: additional_features = additional_features.unsqueeze(dim=0) else: # ['CLS'] token # Kacper: this shouldnt be the case for our project embedded_sentences = embedded_sentences[:, :, 0, :] embedded_sentences = self.dropout(embedded_sentences) batch_size, num_sentences, _ = list(embedded_sentences.size()) sent_mask = (mask.sum(dim=2) != 0) embedded_sentences = self.self_attn(embedded_sentences, sent_mask) if additional_features is not None: embedded_sentences = torch.cat( (embedded_sentences, additional_features), dim=-1) # Kacper: we unwrap the time dimension of a tensor into the 1st dimension (batch), # Kacper: apply a linear layer and wrap the the time dimension back # Kacper: I would suspect it is happening only for embeddings related to the [SEP] tokens label_logits = self.time_distributed_aggregate_feedforward( embedded_sentences) # label_logits: batch_size, num_sentences, num_labels self.track_embedding["logits"] = { "size": list(label_logits.size()), "dim": label_logits.dim() } #print(self.track_embedding) self.track_embedding_list.append(deepcopy(self.track_embedding)) with open(path_json, 'w') as json_out: json.dump(self.track_embedding_list, json_out) if self.labels_are_scores: label_probs = label_logits else: label_probs = torch.nn.functional.softmax(label_logits, dim=-1) # Create output dictionary for the trainer # Compute loss and epoch metrics output_dict = {"action_probs": label_probs} # ===================================================================== if self.with_crf: # Layer 4 = CRF layer across labels of sentences in an abstract mask_sentences = (labels != -1) best_paths = self.crf.viterbi_tags(label_logits, mask_sentences) # # # Just get the tags and ignore the score. predicted_labels = [x for x, y in best_paths] # print(f"len(predicted_labels):{len(predicted_labels)}, (predicted_labels):{predicted_labels}") label_loss = 0.0 if labels is not None: # Compute cross entropy loss # Kacper: reshape logits to be of the following shape in view() flattened_logits = label_logits.view((batch_size * num_sentences), self.num_labels) # Make labels to be contiguous in memory, reshape it so it is in a one dimension flattened_gold = labels.contiguous().view( -1) # Kacper: True labels if not self.with_crf: # Kacper: We are only interested in this part of the code since we don't use crf # Kacper: Get a loss (MSE if sci_sum is True or Crossentropy) label_loss = self.loss(flattened_logits.squeeze(), flattened_gold) if confidences is not None: label_loss = label_loss * confidences.type_as( label_loss).view(-1) label_loss = label_loss.mean() # Kacper: Get a mean loss # Kacper: Get a probabilities from the logits flattened_probs = torch.softmax(flattened_logits, dim=-1) else: # Kacper: We are not interested in this if statement branch (for our project) clamped_labels = torch.clamp(labels, min=0) log_likelihood = self.crf(label_logits, clamped_labels, mask_sentences) label_loss = -log_likelihood # compute categorical accuracy crf_label_probs = label_logits * 0. for i, instance_labels in enumerate(predicted_labels): for j, label_id in enumerate(instance_labels): crf_label_probs[i, j, label_id] = 1 flattened_probs = crf_label_probs.view( (batch_size * num_sentences), self.num_labels) if not self.labels_are_scores: # Kacper: this will be a case for us as well because labels are numerical for Pubmed data evaluation_mask = (flattened_gold != -1) # Kacper: CategoricalAccuracy is computed in this case self.label_accuracy(flattened_probs.float().contiguous(), flattened_gold.squeeze(-1), mask=evaluation_mask) # compute F1 per label for label_index in range(self.num_labels): label_name = self.vocab.get_token_from_index( namespace='labels', index=label_index) metric = self.label_f1_metrics[label_name] metric(flattened_probs, flattened_gold, mask=evaluation_mask) if labels is not None: output_dict["loss"] = label_loss output_dict['action_logits'] = label_logits return output_dict def get_metrics(self, reset: bool = False): # Kacper: this function has to implemented due to API requirements for AllenNLP # Kacper: so it can be run automatically with a config file metric_dict = {} if not self.labels_are_scores: type_accuracy = self.label_accuracy.get_metric(reset) metric_dict['acc'] = type_accuracy average_F1 = 0.0 for name, metric in self.label_f1_metrics.items(): metric_val = metric.get_metric(reset) metric_dict[name + 'F'] = metric_val[2] average_F1 += metric_val[2] average_F1 /= len(self.label_f1_metrics.items()) metric_dict['avgF'] = average_F1 return metric_dict
def train(opts: Settings): # === INSTANTIATE ENVIRONMENT === # gym.make() but with imports configured as specified in arg. _gym_make = partial(gym_make, opts.imports) subproc_gym_make = subproc(_gym_make) # If `opts.subproc==True`, invoke gym.make() in a subprocess, # and treat the resultant instance as a `gym.Env`. make_env = subproc_gym_make if opts.subproc else _gym_make def get_env(index: int): env = make_env(opts.env_id) env.seed(index) env.reset() return env env = MultiEnv(get_env, opts.num_envs) entry_type = [ ('state', env.observation_space.dtype, env.observation_space.shape), ('action', env.action_space.dtype, env.action_space.shape), ('reward', np.float32, (1, )), # ('state1', env.observation_space.dtype, env.observation_space.shape), ('done', np.bool, (1, )), ('value', np.float32, (1, )), ('log_prob', np.float32, env.action_space.shape) ] # === NORMALIZERS FOR INPUTS === reward_normalizer = ExponentialMovingGaussian( alpha=opts.reward_normalizer_alpha) state_normalizer = ExponentialMovingGaussian( alpha=opts.state_normalizer_alpha) # === INSTANTIATE MEMORY === memory = ContiguousRingBuffer(capacity=opts.update_steps, dims=(opts.num_envs, ), dtype=entry_type) # === INSTANTIATE POLICY === # FIXME(ycho): Instead of assuming 1D box spaces, # explicitly wrap envs with flatten()... device = th.device(opts.device) policy = AC(env.observation_space.shape[0], env.action_space.shape[0], opts.ac).to(device) # === INSTANTIATE AGENT === ppo = PPO(policy, memory, device, opts.ppo) # === TRAIN === states = env.reset() dones = np.full((opts.num_envs, 1), False, dtype=np.bool) returns = np.zeros(opts.num_envs, dtype=np.float32) # === LOGGER === # TODO(ycho): Configure logger writer = SummaryWriter() writer.add_graph(policy, th.as_tensor(states).float().to(device)) # === CALLBACKS === save_cb = SaveCallback( opts.save_steps, opts.ckpt_path, lambda: { 'settings': opts, 'state_dict': policy.state_dict(), 'reward_normalizer': reward_normalizer.params(), 'state_normalizer': state_normalizer.params() }) # === VARIABLES FOR DEBUGGING / LOG TRACKING === reset_count = 0 start_time = time.time() # === START TRAINING === step = 0 while step < opts.max_steps: # Reset any env that has reached termination. # FIXME(ycho): assumes isinstance(env, MultiEnv), of course. for i in range(opts.num_envs): if not dones[i]: continue states[i][:] = env.envs[i].reset() returns[i] = 0.0 reset_count += 1 # NOTE(ycho): Workaround for the current limitation of `MultiEnv`. # action = [env.action_space.sample() for _ in range(opts.num_envs)] # sanitize `states` arg. states = np.asarray(states).astype(np.float32) # Add states to stats for normalization. for s in states: state_normalizer.add(s) # Normalize states in-place. states = state_normalizer.normalize(states) states = np.clip(states, -10.0, 10.0) # clip to +-10 stddev with th.no_grad(): action, value, log_prob = ppo.act(states, True) # NOTE(ycho): Clip action within valid domain... clipped_action = np.clip(action, env.action_space.low, env.action_space.high) # Step according to above action. out = env.step(clipped_action) # Format entry. nxt_states, rewards, dones, _ = out # Add rewards to stats for normalization. # returns[np.asarray(dones).reshape(-1).astype(np.bool)] = 0.0 returns = returns * opts.gae.gamma + np.reshape(rewards, -1) # NOTE(ycho): collect stats on `returns` instead of `rewards`. # for r in rewards: # reward_normalizer.add(r) for r in returns: reward_normalizer.add(r) # Train if buffer full ... if memory.is_full: writer.add_scalar('reward_mean', reward_normalizer.mean, global_step=step) writer.add_scalar('reward_var', reward_normalizer.var, global_step=step) writer.add_scalar('log_std', policy.log_std.detach().cpu().numpy()[0], global_step=step) writer.add_scalar('fps', step / (time.time() - start_time), global_step=step) # NOTE(ycho): Don't rely on printed reward stats for tracking # training progress ... use tensorboard instead. print('== step {} =='.format(step)) # Log reward before overwriting with normalized values. print('rew = mean {} min {} max {} std {}'.format( memory['reward'].mean(), memory['reward'].min(), memory['reward'].max(), memory['reward'].std())) # print('rm {} rv {}'.format(reward_normalizer.mean, # reward_normalizer.var)) # NOTE(ycho): States have already been normalized, # since those states were utilized as input for PPO action. # After that, the normalized states were inserted in memory. # memory['state'] = state_normalizer.normalize(memory['state']) # NOTE(ycho): I think it's fine to delay reward normalization to this point. # memory['reward'] = reward_normalizer.normalize(memory['reward']) # NOTE(ycho): maybe the proper thing to do is: # memory['reward'] = (memory['reward'] - reward_normalizer.mean) / np.sqrt(return_normalizer.var) memory['reward'] /= np.sqrt(reward_normalizer.var) memory['reward'] = np.clip(memory['reward'], -10.0, 10.0) # Create training data slices from memory ... dones = np.asarray(dones).reshape(opts.num_envs, 1) advs, rets = gae(memory, value, dones, opts.gae) # print('std = {}'.format(ppo.policy.log_std.exp())) ucount = 0 info = None for _ in range(opts.num_epochs): for i in range(0, len(memory), opts.batch_size): # Prepare current minibatch dataset ... exp = memory[i:i + opts.batch_size] act = exp['action'] obs = exp['state'] old_lp = exp['log_prob'] # old_v = exp['value'] # NOTE(ycho): unused adv = advs[i:i + opts.batch_size] ret = rets[i:i + opts.batch_size] # Evaluate what had been done ... # NOTE(ycho): wouldn't new_v == old_v # and new_lp == old_lp for the very first one in the batch?? # hmm .... new_v, new_lp, entropy = ppo.evaluate( obs.copy(), act.copy()) info_i = {} loss = ppo.compute_loss(obs.copy(), act.copy(), old_lp.copy(), new_v, new_lp, entropy, adv, ret, info_i) # NOTE(ycho): Below, only required for logging if True: with th.no_grad(): if info is None: info = info_i else: for k in info.keys(): info[k] += info_i[k] ucount += 1 # Optimization step ppo.optimizer.zero_grad() loss.backward() # Clip grad norm th.nn.utils.clip_grad_norm_(ppo.policy.parameters(), opts.ppo.max_grad_norm) ppo.optimizer.step() for k, v in info.items(): writer.add_scalar(k, v.detach().cpu().numpy() / ucount, global_step=step) # Empty the memory ! memory.reset() # Append to memory. entry = list( zip(*( states, action, rewards, # nxt_states, dones, value, log_prob))) memory.append(entry) # Cache `states`, update steps and continue. states = nxt_states step += opts.num_envs save_cb.on_step(step) writer.close() # Save ... th.save( { 'settings': opts, 'state_dict': policy.state_dict(), 'reward_normalizer': reward_normalizer.params(), 'state_normalizer': state_normalizer.params() }, opts.model_path)
def train(load_path=None): # Define exp name exp_name = input("Experience Name (You may leave this blank): ") exp_name_ = datetime.now().strftime("%m-%d-%H-%M") + ( '-' + exp_name) if exp_name else '' # Prepare status save location if not os.path.exists(save_base_path): os.mkdir(save_base_path) if not os.path.exists(os.path.join(save_base_path, identifier)): os.mkdir(os.path.join(save_base_path, identifier)) save_path = os.path.join(save_base_path, identifier, exp_name_) if not os.path.exists(save_path): os.mkdir(save_path) # Save constants.py as json file with open(os.path.join(save_path, 'config.json'), 'w') as f: json.dump(C, f, default=lambda x: str(x), indent=True) # Load dipole kernel d = load_dipole().to(device) # Prepare training and validation datasets train_dataset = QSMDataset('train') b_mean = train_dataset.X_mean b_std = train_dataset.X_std y_mean = train_dataset.Y_mean y_std = train_dataset.Y_std train_loader = DataLoader(train_dataset, num_workers=4, batch_size=batch_size, shuffle=True, drop_last=True) val_loader = DataLoader(QSMDataset('val'), batch_size=1, shuffle=False) # Initialize training elements if load_path: checkpoint = torch.load(load_path, map_location='cpu') start_epoch = checkpoint.get('epoch') network = checkpoint.get('model') optimizer = checkpoint.get('optimizer') del checkpoint network = network.to(device) for state in optimizer.state.values(): for k, v in state.items(): if torch.is_tensor(v): state[k] = v.to(device) else: start_epoch = 0 network = QSMNet().to(device) optimizer = optim.Adam(network.parameters(), lr=learning_rate) # Initialize tensorboard logging writer = SummaryWriter(comment='_' + identifier + '_' + exp_name) writer.add_graph(network, torch.zeros(1, 1, 64, 64, 64, device=device)) # Announce training print('{} training network "{}" for {} epoches on {}{}.'.format( 'Resume' if load_path else 'Begin', identifier, train_epochs, device.type, ':' + str(device.index) if device.index is not None else '')) # Begin training step = start_epoch * len(train_loader) # Global step count for tensorboard for epoch in range(start_epoch, start_epoch + train_epochs): network.train() loss_history = [] if epoch < one_to_two - 1: print("Loss gear phase one") total_loss.phase = 'one' else: print("Loss gear phase two") total_loss.phase = 'two' for batch, (b, y, m) in enumerate(train_loader): b, y, m = b.to(device), y.to(device), m.to(device) optimizer.zero_grad() chi = network(b) loss, loss_parts = total_loss(chi, y, b, d, m, b_mean, b_std, y_mean, y_std) loss.backward() step += 1 optimizer.step() loss_cpu = loss.detach().cpu().item() loss_history.append(loss_cpu) torch.cuda.empty_cache() with torch.no_grad(): # Track loss with tensorboard writer.add_scalar('Loss/Total loss', loss_cpu, step) for loss_name, loss_value in loss_parts.items(): writer.add_scalar('Loss/' + loss_name, loss_value, step) # Print step status if (batch + 1) % print_every == 0: train_log = 'Epoch {:2d}/{:2d}\tLoss: {:.6f}\tTrain: [{}/{} ({:.0f}%)] '.format( epoch + 1, start_epoch + train_epochs, loss_cpu, batch + 1, len(train_loader), 100. * batch / len(train_loader)) print(train_log, end='\r') sys.stdout.flush() with torch.no_grad(): # Print epoch status last_epoch = epoch + 1 average_loss = sum(loss_history) / len(loss_history) epoch_end_log = "Epoch {:02d} completed, Average Loss is {:.6f}. ".format( last_epoch, average_loss) print(epoch_end_log) # Save training status if last_epoch % save_every == 0: ckpt_save_path = os.path.join( save_path, '{}-epoch{:02d}.ckpt'.format(identifier, last_epoch)) torch.save( { 'epoch': last_epoch, 'model': network, 'optimizer': optimizer, }, ckpt_save_path) print("Saved Model on epoch {}.".format(last_epoch)) # Evaluate model #evaluate_model('validation', network, val_loader, last_epoch, save_path, writer) challenge_evaluation(exp_name_, network, writer, last_epoch) writer.close()
class LMTrainer: def __init__(self, model, dataset, checkpoint_filename): self.checkpoint_filename = checkpoint_filename self.model = model.cuda() self.dataset = dataset self.writer = SummaryWriter() self.global_step = 0 self.last_log = datetime.datetime.now() self.last_valid = datetime.datetime.now() # Dataset train_dataset, test_dataset = random_split(dataset, [90, 10]) self.loaders = { "train": DataLoader( train_dataset, batch_size=config.BATCH_SIZE, shuffle=True, drop_last=True, ), "test": DataLoader(test_dataset, batch_size=config.BATCH_SIZE, shuffle=True, drop_last=True), } inputs = next(iter(self.loaders["test"])) model(inputs.cuda()) self.writer.add_graph(model, inputs.cuda()) # Optimizer self.optim = torch.optim.lr_scheduler.CyclicLR( torch.optim.SGD(self.model.parameters(), lr=config.LEARNING_RATE, momentum=0.9), base_lr=0, max_lr=config.LEARNING_RATE, ) self.best_valid = 1e9 self.patience = 0 def run_batch_learn(self, i, batch): self.optim.optimizer.zero_grad() self.model.train() loss, metrics = self.compute_loss_and_metrics_from_indices(batch) loss.backward() self.optim.optimizer.step() self.optim.step() self.writer.add_scalar( "train/lr", torch.Tensor([self.optim.get_lr()]), global_step=self.global_step, ) self.global_step += config.BATCH_SIZE def evaluate_batch(self, mode, batch): loss, metrics = self.compute_loss_and_metrics_from_indices(batch) for name, metric in sorted(metrics.items()): scalar_name = f"{mode}/{name}" self.writer.add_scalar(scalar_name, metric, global_step=self.global_step) return loss def evaluate(self, mode): total_loss = 0 with torch.no_grad(): self.model.eval() loader = self.loaders[mode] for batch in tqdm.tqdm(loader, total=len(loader), desc=mode): batch = batch.cuda() loss = self.evaluate_batch(mode, batch) total_loss += loss.item() return total_loss def run_epoch(self): loader = self.loaders["train"] for i, batch in tqdm.tqdm(enumerate(loader), total=len(loader), desc="train"): batch = batch.cuda() self.run_batch_learn(i, batch) now = datetime.datetime.now() if now - self.last_log > config.LOG_INTERVAL: self.last_log = now self.evaluate_batch("train", batch) if now - self.last_valid > config.TEST_INTERVAL: valid_loss = self.evaluate("test") self.last_valid = now if valid_loss < self.best_valid: torch.save(self.model, self.checkpoint_filename) self.best_valid = valid_loss self.patience = 0 else: self.patience += 1 msg = f"Did not improve valid loss, model not saved (patience: {self.patience})" logging.info(msg) def compute_loss_and_metrics_from_indices(self, indices): attn = self.model.forward(indices) # attn: [B, C, D] decoded = self.model.decoder(attn) # decoded: [B, C, V] B, C, V = decoded.shape x = (decoded[:, :-1, :].contiguous().view(-1, decoded.size(-1)) ) # [B * (C - 1), V] logits = F.log_softmax(x, dim=-1) # [B * (C - 1), V] next_words = indices[:, 1:].contiguous().view(-1) # [B * (C - 1)] prediction_loss = F.nll_loss(logits, next_words, reduction="sum") / attn.size(0) loss_words = prediction_loss / config.CONTEXT_SIZE _, topk = logits.topk(10, -1) topks = { f"P_{k}": (topk[:, :k] == next_words.view(-1, 1)).sum(-1).float().mean() * 100 for k in [1, 5, 10] } losses = { "loss/context": prediction_loss, "loss/words": loss_words, **topks } loss = prediction_loss if random.random() < 0.01: plt.clf() values = torch.exp(logits.reshape(B, C - 1, V)[0, -1, :]) plt.plot(values.detach().cpu().numpy()) mu = 500 sigma = 100 x = np.linspace(mu - 5 * sigma, mu + 5 * sigma, 100) plt.plot(x, scipy.stats.norm.pdf(x, mu, sigma)) mu = 500 sigma = 10 x = np.linspace(mu - 5 * sigma, mu + 5 * sigma, 100) plt.plot(x, scipy.stats.norm.pdf(x, mu, sigma)) global INDEX filename = f"figures/out{INDEX}.png" plt.savefig(filename) INDEX += 1 print(f"Saved {filename}") return loss, losses