def train( net: torch.nn.Module, dataloader: DataLoader, criterion, # torch.nn._Loss optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler._LRScheduler, epochs=10, ) -> Iterable[torch.nn.Module]: net.to(TORCH_DEVICE, dtype=torch.float) i = 0 for epoch in range(epochs): # loop over the dataset multiple times for j, (inputs, labels) in enumerate(dataloader): # get the inputs inputs, labels = ( inputs.to(TORCH_DEVICE, dtype=torch.float), labels.to(TORCH_DEVICE, dtype=torch.long) ) i += len(inputs) # zero the parameter gradients optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() if np.isnan(loss.item()): raise Exception('gradients blew up') writer.add_scalar('Train/Loss', loss, i) writer.add_scalar('Train/LR', optimizer.param_groups[0]['lr'], i) scheduler.step() yield net, loss.item()
def load_checkpoint( checkpoint_file: Union[str, Path], model: torch.nn.Module, optimizer: torch.optim.Optimizer = None, scheduler: torch.optim.lr_scheduler._LRScheduler = None, map_location: Union[torch.device, str, Mapping[str, str], Callable] = None, ) -> None: """Shortcut for loading checkpoint state. Args: checkpoint_file (Union[str, Path]): path to checkpoint model (torch.nn.Module): model to initialize with checkpoint weights optimizer (torch.optim.Optimizer, optional): optimizer to initialize with checkpoint weights. If `None` then will be ignored. Default is None. scheduler (torch.optim.lr_scheduler._LRScheduler, optional): scheduler to initialize with checkpoint weights. If `None` then will be ignored. Default is None. map_location (Union[torch.device, str, Mapping[int, str], Callable], optional): location to use for loading checkpoint content. More about possible locations - https://pytorch.org/docs/master/generated/torch.load.html Default is None. """ checkpoint = torch.load(str(checkpoint_file), map_location=map_location) loaded_items = [] if "model_state_dict" in checkpoint and model is not None: state_dict = checkpoint["model_state_dict"] if isinstance(model, torch.nn.DataParallel): model.module.load_state_dict(state_dict) else: model.load_state_dict(state_dict) loaded_items.append("model") if "optimizer_state_dict" in checkpoint and optimizer is not None: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) loaded_items.append("optimizer") if "scheduler_state_dict" in checkpoint and scheduler is not None: scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) loaded_items.append("scheduler") if loaded_items: print("<= Loaded {} from '{}'".format(", ".join(loaded_items), checkpoint_file)) if "stage" in checkpoint: print("Stage -", checkpoint["stage"]) if "epoch" in checkpoint: print("Epoch -", checkpoint["epoch"]) if "metrics" in checkpoint: print("Metrics:") pprint(checkpoint["metrics"])
def train( model: torch.nn.Module, train_data: torch.utils.data.Dataset, val_data: torch.utils.data.Dataset, test_data: torch.utils.data.Dataset, optimiser: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler._LRScheduler, criterion: Optional[Callable[..., torch.Tensor]], loggers: Iterable[Logger], config: argparse.Namespace, ): """Train a model.""" dataloader = DataLoader(train_data, batch_size=config.batch_size, shuffle=True, num_workers=0) class_count = len(set(dataloader.dataset.tensors[1].tolist())) for epoch in range(config.epochs): model.train() for batch, (feats, labels) in enumerate(dataloader): # Move data to GPU feats = feats.to(config.device) # Convert labels to one-hots if using BCEWithLogitsLoss if isinstance(criterion, torch.nn.BCEWithLogitsLoss): labels = F.one_hot(labels, num_classes=class_count).type(torch.float32) labels = labels.to(config.device) optimiser.zero_grad() clean_posteriors, clean_activations = model(feats) # TODO Make sure we apply log for NONE case (in new loss class) loss = criterion(clean_posteriors, labels) loss.backward() optimiser.step() preds = torch.argmax(clean_posteriors, dim=-1).cpu().numpy() if batch % config.log_step == config.log_step - 1 or batch == len( dataloader) - 1: metrics = { "train/epoch": epoch + batch * dataloader.batch_size / len(dataloader.dataset), "train/lr": scheduler.get_last_lr()[0], "train/loss": loss.item(), "train/accuracy": accuracy_score(labels.cpu().numpy(), preds), } for logger in loggers: logger(metrics) evaluate(model, epoch + 1, val_data, criterion, loggers, config) test(model, test_data, loggers, config) scheduler.step()
def train(train_loader: DataLoader, val_loader: DataLoader, model: nn.Module, criterion: nn.Module, optimizer: Optimizer, scheduler: torch.optim.lr_scheduler._LRScheduler, args): """ 训练模型 :param train_loader: 训练集 :param val_loader: 验证集 :param model: 模型 :param criterion: 损失函数 :param optimizer: 优化器 :param args: 训练超参 """ writer = SummaryWriter(args.logdir) best_val_acc1 = 0 learning_rate = 0 for epoch in range(args.epochs): if args.distributed: train_loader.sampler.set_epoch(epoch) learning_rate = scheduler.get_last_lr() if isinstance(learning_rate, list): learning_rate = learning_rate[0] # 训练一个epoch,并在验证集上评估 train_loss, train_acc1 = train_epoch(train_loader, model, criterion, optimizer, epoch, args) val_acc1, val_loss, _ = test(val_loader, model, criterion, args) scheduler.step() # 保存当前及最好的acc@1的checkpoint is_best = val_acc1 > best_val_acc1 best_val_acc1 = max(val_acc1, best_val_acc1) save_checkpoint({ # 'epoch': epoch + 1, # 'arch': args.arch, 'state_dict': model.module.state_dict(), # 'best_acc1': best_val_acc1, # 'optimizer': optimizer.state_dict(), }, is_best, args) writer.add_scalar('learning rate', learning_rate, epoch) writer.add_scalar('Train/Loss', train_loss, epoch) writer.add_scalar('Train/Accuracy', train_acc1, epoch) writer.add_scalar('Val/Loss', val_loss, epoch) writer.add_scalar('Val/Accuracy', val_acc1, epoch) writer.flush() writer.close() logging.info(f'Training Over with lr={learning_rate}~~')
def fit( # noqa: C901 self, train_generator: DataLoader, val_generator: DataLoader, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler._LRScheduler, epochs: int, criterion_dict: dict[str, nn.Module], normalizer_dict: dict[str, Normalizer], model_name: str, run_id: int, checkpoint: bool = True, writer: SummaryWriter = None, verbose: bool = True, patience: int = None, ) -> None: """Convenience class to carry out training loop. Args: train_generator (DataLoader): Dataloader containing training data. val_generator (DataLoader): Dataloader containing validation data. optimizer (torch.optim.Optimizer): Optimizer used to carry out parameter updates. scheduler (torch.optim.lr_scheduler._LRScheduler): Scheduler used to adjust Optimizer during training. epochs (int): Number of epochs to train for. criterion_dict (dict[str, nn.Module]): Dictionary of losses to apply for each task. normalizer_dict (dict[str, Normalizer]): Dictionary of Normalizers to apply to each task. model_name (str): String describing the model. run_id (int): Unique identifier of the model run. checkpoint (bool, optional): Whether to save model checkpoints. Defaults to True. writer (SummaryWriter, optional): TensorBoard writer to save logs in. Defaults to None. verbose (bool, optional): Whether to print out intermediate results. Defaults to True. patience (int, optional): Patience for early stopping. Defaults to None. """ start_epoch = self.epoch try: for epoch in range(start_epoch, start_epoch + epochs): self.epoch += 1 # Training t_metrics = self.evaluate( generator=train_generator, criterion_dict=criterion_dict, optimizer=optimizer, normalizer_dict=normalizer_dict, action="train", verbose=verbose, ) if writer is not None: for task, metrics in t_metrics.items(): for metric, val in metrics.items(): writer.add_scalar(f"{task}/train/{metric}", val, epoch) if verbose: print(f"Epoch: [{epoch}/{start_epoch + epochs - 1}]") for task, metrics in t_metrics.items(): metrics_str = "".join([ f"{key} {val:.3f}\t" for key, val in metrics.items() ]) print(f"Train \t\t: {task} - {metrics_str}") # Validation if val_generator is not None: with torch.no_grad(): # evaluate on validation set v_metrics = self.evaluate( generator=val_generator, criterion_dict=criterion_dict, optimizer=None, normalizer_dict=normalizer_dict, action="val", ) if writer is not None: for task, metrics in v_metrics.items(): for metric, val in metrics.items(): writer.add_scalar( f"{task}/validation/{metric}", val, epoch) if verbose: for task, metrics in v_metrics.items(): metrics_str = "".join([ f"{key} {val:.3f}\t" for key, val in metrics.items() ]) print(f"Validation \t: {task} - {metrics_str}") # TODO test all tasks to see if they are best, # save a best model if any is best. # TODO what are the costs of this approach. # It could involve saving a lot of models? is_best: list[bool] = [] for name in self.best_val_scores: if self.task_dict[name] == "regression": if v_metrics[name]["MAE"] < self.best_val_scores[ name]: self.best_val_scores[name] = v_metrics[name][ "MAE"] is_best.append(True) is_best.append(False) elif self.task_dict[name] == "classification": if v_metrics[name]["Acc"] > self.best_val_scores[ name]: self.best_val_scores[name] = v_metrics[name][ "Acc"] is_best.append(True) is_best.append(False) if any(is_best): self.es_patience = 0 else: self.es_patience += 1 if patience and self.es_patience > patience: print( "Stopping early due to lack of improvement on Validation set" ) break if checkpoint: checkpoint_dict = { "model_params": self.model_params, "state_dict": self.state_dict(), "epoch": self.epoch, "best_val_score": self.best_val_scores, "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "normalizer_dict": { task: n.state_dict() if isinstance(n, Normalizer) else None for task, n in normalizer_dict.items() }, } # TODO saving a model at each epoch may be slow? save_checkpoint(checkpoint_dict, False, model_name, run_id) # TODO when to save best models? should this be done task-wise in # the multi-task case? # save_checkpoint(checkpoint_dict, is_best, model_name, run_id) scheduler.step() # catch memory leak gc.collect() except KeyboardInterrupt: pass if writer is not None: writer.close()
def train_model( model: torch.nn.Module, train_dataset: Dataset, val_dataset: Dataset, val_transforms: Callable, loss_fn: Callable, metrics: Dict[str, Callable], optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler._LRScheduler = None, loader_workers: int = 4, batch_size: int = 32, n_epochs: int = 20, device: str = 'cuda:0', logger: ExperimentLogger = None, snapshot_dir: str = './snapshots', cycle: int = None, ): # TODO: adaptive learning rate, early stopping conditions train_loader = DataLoader(train_dataset, batch_size, shuffle=True, num_workers=loader_workers) val_loader = DataLoader(val_dataset, batch_size, shuffle=False, num_workers=loader_workers) if cycle is not None: snapshot_dir = os.path.join(snapshot_dir, f'cycle_{cycle}') os.makedirs(snapshot_dir) best_mertics_dict = {metric_name: 0 for metric_name in metrics} best_val_loss = 1e10 for epoch in range(n_epochs): metrics_dict = defaultdict(list) train_loss = [] val_loss = [] model.train(True) for batch in train_loader: image, mask = batch['image'], batch['mask'] image = image.to(device) y_pred = model(image) loss = loss_fn(y_pred, mask.to(device)) for metric_name, metric in metrics.items(): metrics_dict[f'train_{metric_name}'] += metric( y_pred, mask.to(device)) loss.backward() optimizer.step() optimizer.zero_grad() train_loss.append(loss.item()) model.eval() orig_masks = [] pred_masks = [] e_loss = [] for batch in val_loader: image, mask_raw = batch['image'], batch['mask'] image = image.to(device) with torch.no_grad(): y_pred_raw = model(image) y_pred = torch.sigmoid(y_pred_raw).cpu().squeeze().numpy() mask = mask_raw.squeeze().numpy() e_loss.append(loss_fn(y_pred_raw, mask_raw.to(device))) for i in range(mask.shape[0]): predicted_mask = y_pred original_mask = mask pred_masks.append(val_transforms(predicted_mask.squeeze())) orig_masks.append(val_transforms(original_mask.squeeze())) pred_masks_numpy = np.array(pred_masks) orig_masks = torch.Tensor(orig_masks).unsqueeze(1) pred_masks = torch.Tensor(pred_masks_numpy).unsqueeze(1) val_loss.append(np.mean(e_loss)) for metric_name, metric in metrics.items(): metrics_dict[f'val_{metric_name}'] += metric(pred_masks, orig_masks, th=0.5) if scheduler: scheduler.step(epoch=epoch) logger.log( epoch=epoch, train_loss=np.mean(train_loss), train_iou=np.mean(metrics_dict['train_iou']), train_iout=np.mean(metrics_dict['train_iout']), val_loss=np.mean(val_loss), val_iou=np.mean(metrics_dict['val_iou']), val_iout=np.mean(metrics_dict['val_iout']), lr=optimizer.param_groups[0]['lr'], ) logger.log_val_metrics(','.join( [f'{x:.4f}' for x in metrics_dict['val_iou']])) if np.mean(val_loss) < best_val_loss: best_val_loss = np.mean(val_loss) snapshot_name = f'epoch_{epoch}.loss_{np.mean(val_loss):.4f}.iout_{np.mean(metrics_dict["val_iout"]):.3f}' np.save(arr=pred_masks_numpy, file=os.path.join(snapshot_dir, 'val_masks_' + snapshot_name)) torch.save(model.state_dict(), os.path.join(snapshot_dir, snapshot_name)) for metric_name in metrics: if best_mertics_dict[metric_name] < np.mean( metrics_dict[f'val_{metric_name}']): best_mertics_dict[metric_name] = np.mean( metrics_dict[f'val_{metric_name}']) snapshot_name = f'epoch_{epoch}.loss_{np.mean(val_loss):.4f}.iout_{np.mean(metrics_dict["val_iout"]):.3f}' np.save(arr=pred_masks_numpy, file=os.path.join(snapshot_dir, 'val_masks_' + snapshot_name)) torch.save(model.state_dict(), os.path.join(snapshot_dir, snapshot_name))
def attack(self, epochs: int, optimizer: torch.optim.Optimizer, lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None, validate_interval: int = 1, save: bool = False, verbose: bool = True, **kwargs): if verbose: print('train mask generator') self.mark_generator.requires_grad_(False) self.mask_generator.requires_grad_() self.model.requires_grad_(False) self.train_mask_generator(verbose=verbose) if verbose: print() print('train mark generator and model') self.mark_generator.requires_grad_() self.mask_generator.requires_grad_(False) if not self.natural: params: list[nn.Parameter] = [] for param_group in optimizer.param_groups: params.extend(param_group['params']) self.model.activate_params(params) mark_optimizer = torch.optim.Adam(self.mark_generator.parameters(), lr=1e-2, betas=(0.5, 0.9)) mark_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( mark_optimizer, T_max=epochs) loader = self.dataset.loader['train'] dataset = loader.dataset logger = MetricLogger() logger.create_meters(loss=None, div=None, ce=None) if validate_interval != 0: best_validate_result = self.validate_fn(verbose=verbose) best_asr = best_validate_result[0] for _epoch in range(epochs): _epoch += 1 idx = torch.randperm(len(dataset)) pos = 0 logger.reset() if not self.natural: self.model.train() self.mark_generator.train() header: str = '{blue_light}{0}: {1}{reset}'.format( 'Epoch', output_iter(_epoch, epochs), **ansi) header = header.ljust(max(len('Epoch'), 30) + get_ansi_len(header)) for data in logger.log_every(loader, header=header) if verbose else loader: if not self.natural: optimizer.zero_grad() mark_optimizer.zero_grad() _input, _label = self.model.get_data(data) batch_size = len(_input) data2 = sample_batch(dataset, idx=idx[pos:pos + batch_size]) _input2, _label2 = self.model.get_data(data2) pos += batch_size final_input, final_label = _input.clone(), _label.clone() # generate trigger input trigger_dec, trigger_int = math.modf(len(_label) * self.poison_percent) trigger_int = int(trigger_int) if random.uniform(0, 1) < trigger_dec: trigger_int += 1 x = _input[:trigger_int] trigger_mark, trigger_mask = self.get_mark(x), self.get_mask(x) trigger_input = x + trigger_mask * (trigger_mark - x) final_input[:trigger_int] = trigger_input final_label[:trigger_int] = self.target_class # generate cross input cross_dec, cross_int = math.modf(len(_label) * self.cross_percent) cross_int = int(cross_int) if random.uniform(0, 1) < cross_dec: cross_int += 1 x = _input[trigger_int:trigger_int + cross_int] x2 = _input2[trigger_int:trigger_int + cross_int] cross_mark, cross_mask = self.get_mark(x2), self.get_mask(x2) cross_input = x + cross_mask * (cross_mark - x) final_input[trigger_int:trigger_int + cross_int] = cross_input # div loss if len(trigger_input) <= len(cross_input): length = len(trigger_input) cross_input = cross_input[:length] cross_mark = cross_mark[:length] cross_mask = cross_mask[:length] else: length = len(cross_input) trigger_input = trigger_input[:length] trigger_mark = trigger_mark[:length] trigger_mask = trigger_mask[:length] input_dist: torch.Tensor = (trigger_input - cross_input).flatten(1).norm(p=2, dim=1) mark_dist: torch.Tensor = (trigger_mark - cross_mark).flatten(1).norm(p=2, dim=1) + 1e-5 loss_ce = self.model.loss(final_input, final_label) loss_div = input_dist.div(mark_dist).mean() loss = loss_ce + self.lambda_div * loss_div loss.backward() if not self.natural: optimizer.step() mark_optimizer.step() logger.update(n=batch_size, loss=loss.item(), div=loss_div.item(), ce=loss_ce.item()) if not self.natural and lr_scheduler: lr_scheduler.step() mark_scheduler.step() if not self.natural: self.model.eval() self.mark_generator.eval() if validate_interval != 0 and (_epoch % validate_interval == 0 or _epoch == epochs): validate_result = self.validate_fn(verbose=verbose) cur_asr = validate_result[0] if cur_asr >= best_asr: best_validate_result = validate_result best_asr = cur_asr if save: self.save() if not self.natural: optimizer.zero_grad() mark_optimizer.zero_grad() self.mark_generator.requires_grad_(False) self.mask_generator.requires_grad_(False) self.model.requires_grad_(False) return best_validate_result
def train_epoch(train_loader: torch.utils.data.DataLoader, base_model: torch.nn.Module, classification_layer: torch.nn.Module, forg_layer: torch.nn.Module, epoch: int, optimizer: torch.optim.Optimizer, lr_scheduler: torch.optim.lr_scheduler._LRScheduler, callback: Optional[VisdomLogger], device: torch.device, args: Any): """ Trains the network for one epoch Parameters ---------- train_loader: torch.utils.data.DataLoader Iterable that loads the training set (x, y) tuples base_model: torch.nn.Module The model architecture that "extract features" from signatures classification_layer: torch.nn.Module The classification layer (from features to predictions of which user wrote the signature) forg_layer: torch.nn.Module The forgery prediction layer (from features to predictions of whether the signature is a forgery). Only used in args.forg = True epoch: int The current epoch (used for reporting) optimizer: torch.optim.Optimizer The optimizer (already initialized) lr_scheduler: torch.optim.lr_scheduler._LRScheduler The learning rate scheduler callback: VisdomLogger (optional) A callback to report the training progress device: torch.device The device (CPU or GPU) to use for training args: Namespace Extra arguments used for training: args.forg: bool Whether forgeries are being used for training args.lamb: float The weight used for the forgery loss (training with forgeries only) Returns ------- None """ step = 0 n_steps = len(train_loader) for batch in train_loader: x, y = batch[0], batch[1] x = torch.tensor(x, dtype=torch.float).to(device) y = torch.tensor(y, dtype=torch.long).to(device) yforg = torch.tensor(batch[2], dtype=torch.float).to(device) # Forward propagation features = base_model(x) if args.forg: if args.loss_type == 'L1': # Eq (3) in https://arxiv.org/abs/1705.05787 logits = classification_layer(features) class_loss = F.cross_entropy(logits, y) forg_logits = forg_layer(features).squeeze() forg_loss = F.binary_cross_entropy_with_logits(forg_logits, yforg) loss = (1 - args.lamb) * class_loss loss += args.lamb * forg_loss else: # Eq (4) in https://arxiv.org/abs/1705.05787 logits = classification_layer(features[yforg == 0]) class_loss = F.cross_entropy(logits, y[yforg == 0]) forg_logits = forg_layer(features).squeeze() forg_loss = F.binary_cross_entropy_with_logits(forg_logits, yforg) loss = (1 - args.lamb) * class_loss loss += args.lamb * forg_loss else: # Eq (1) in https://arxiv.org/abs/1705.05787 logits = classification_layer(features) loss = class_loss = F.cross_entropy(logits, y) # Back propagation optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_value_(optimizer.param_groups[0]['params'], 10) # Update weights optimizer.step() # Logging if callback and step % 100 == 0: iteration = epoch + (step / n_steps) callback.scalar('class_loss', iteration, class_loss.detach()) pred = logits.argmax(1) acc = y[yforg == 0].eq(pred).float().mean() callback.scalar('train_acc', epoch + (step / n_steps), acc.detach()) if args.forg: forg_pred = forg_logits > 0 forg_acc = yforg.long().eq(forg_pred.long()).float().mean() callback.scalar('forg_loss', iteration, forg_loss.detach()) callback.scalar('forg_acc', iteration, forg_acc.detach()) step += 1 lr_scheduler.step()
def train(train_loader: DataLoader, val_loader: DataLoader, model: nn.Module, criterion: nn.Module, optimizer: Optimizer, scheduler: torch.optim.lr_scheduler._LRScheduler, args): """ 训练模型 :param train_loader: 训练集 :param val_loader: 验证集 :param model: 模型 :param criterion: 损失函数 :param optimizer: 优化器 :param args: 训练超参 """ writer = SummaryWriter(args.logdir) # writer.add_graph(model, (torch.rand(1, 3, args.image_size[0], args.image_size[1]),)) global mix_up, multi_scale, bn_gammas, net_weights if mix_up is None: mix_up = MixUp(args) if args.multi_scale and multi_scale is None: multi_scale = MultiScale(args.image_size) if bn_gammas is None: bn_gammas = [ m.weight for m in model.modules() if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.SyncBatchNorm) or isinstance(m, apex.parallel.SyncBatchNorm) ] if net_weights is None: net_weights = [ param for name, param in model.named_parameters() if name[-4:] != 'bias' ] best_val_acc1 = 0 learning_rate = 0 for epoch in range(args.epochs): if args.distributed: train_loader.sampler.set_epoch(epoch) learning_rate = scheduler.get_last_lr() if isinstance(learning_rate, list): learning_rate = learning_rate[0] # 训练一个epoch,并在验证集上评估 train_loss, train_acc1 = train_epoch(train_loader, model, criterion, optimizer, epoch, args) val_acc1, val_loss, _ = test(val_loader, model, criterion, args, is_confuse_matrix=False) scheduler.step() # 保存当前及最好的acc@1的checkpoint is_best = val_acc1 >= best_val_acc1 best_val_acc1 = max(val_acc1, best_val_acc1) save_checkpoint( { # 'epoch': epoch + 1, # 'arch': args.arch, 'state_dict': model.module.state_dict(), # 'best_acc1': best_val_acc1, # 'optimizer': optimizer.state_dict(), }, is_best, args) all_bn_weight = [] for gamma in bn_gammas: all_bn_weight.append(gamma.cpu().data.numpy()) writer.add_histogram('BN gamma', np.concatenate(all_bn_weight, axis=0), epoch) # writer.add_scalars('Loss', {'Train': train_loss, 'Val': val_loss}, epoch) # writer.add_scalars('Accuracy', {'Train': train_acc1, 'Val': val_acc1}, epoch) writer.add_scalar('Train/Loss', train_loss, epoch) writer.add_scalar('Train/Accuracy', train_acc1, epoch) writer.add_scalar('Val/Loss', val_loss, epoch) writer.add_scalar('Val/Accuracy', val_acc1, epoch) writer.add_scalar('learning rate', learning_rate, epoch) writer.flush() writer.close() logging.info(f'Training Over with lr={learning_rate}~~')
def actor_critic_train_per_episode( policy: SimplePolicyContinuous, critic: SimpleCritic, env: gym.Env, optimizer: Optimizer, run_params: RunParams, lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None): """ Trains the actor critic on the given environment. Training is done at the end of each episode, instead of at the end of each step of an episode. This means the agent trains much more frequently. Both discrete and continuous actions spaces are supported. Several features can be optionally enabled: 1) Scaling / normalizing the states / observations 2) Logging training statistics on Tensorboard 3) Render the environment periodically (pick render_frequency in the RunParams) 4) Using a learning rate scheduler """ training_info = TrainingInfo(GAMMA=run_params.gamma) print( f"The goal is a running reward of at least {env.spec.reward_threshold}." ) # https://medium.com/@asteinbach/actor-critic-using-deep-rl-continuous-mountain-car-in-tensorflow-4c1fb2110f7c # says it's crucial to scale the state if run_params.should_scale_states: scaler = setup_observation_scaler(env) writer = run_params.get_tensorboard_writer( env) if run_params.use_tensorboard else None for episode_number in itertools.count( ): # itertools.count() is basically range(+infinity) state = env.reset() # Do a whole episode (upto 10000 steps, don't want infinite steps) for t in range(env.spec.max_episode_steps): if run_params.should_scale_states: state = scale_state(scaler, state) if run_params.continuous_actions: action = select_action_continuous(state, policy, training_info, env) else: action = select_action_discrete(state, policy, training_info) state_value = get_state_value(state, critic) new_state, reward, done, _ = env.step(action) if run_params.should_render(episode_number): env.render() training_info.record_step( state, action, reward, state_value) # Store reward and updates the running reward state = new_state if done: break training_info.update_running_reward() # Add some logging log_on_console(env, episode_number, reward, run_params, t, training_info) log_on_tensorboard(env, episode_number, reward, run_params, t, training_info, writer) # Check if we have solved the environment reliably if env.spec.reward_threshold is not None and training_info.running_reward > env.spec.reward_threshold: print( f"Solved! The running reward is {training_info.running_reward:.2f}, which is above the threshold of " f"{env.spec.reward_threshold}. The last episode ran for {t} steps." ) break train_policy_on_episode(optimizer, training_info, episode_number) if lr_scheduler: lr_scheduler.step(episode_number) close_tensorboard(run_params, writer)