train_data = get_train_data( data_cfg['data_path'], data_cfg['train_data'], vocab, [preprocess_cfg['pa_max_sent_len'], preprocess_cfg['qu_max_sent_len']]) dev_data = get_dev_data( data_cfg['data_path'], data_cfg['dev_data'], vocab, [preprocess_cfg['pa_max_sent_len'], preprocess_cfg['qu_max_sent_len']]) model = RecurrentModel(preprocess_cfg['vocab_size'], model_cfg['embedding_dim'], model_cfg['hidden_size']).to(device) model.train() print(model) optimizer = Adam(model.parameters(), train_cfg["lr"], weight_decay=train_cfg["weight_decay"]) lr_now = train_cfg["lr"] crossentropyloss = torch.nn.CrossEntropyLoss() train_loader = DataLoader(train_data, batch_size=train_cfg["batch_size"], shuffle=True, collate_fn=collate_fn) dev_loader = DataLoader(dev_data, batch_size=dev_cfg["batch_size"], shuffle=True, collate_fn=collate_fn) logging_step = train_cfg["logging_step"]
def train( epochs: int, models_dir: Path, x_cities: List[CityData], y_city: List[CityData], mask_dir: Path, ): model = UNet11().cuda() optimizer = Adam(model.parameters(), lr=3e-4) scheduler = ReduceLROnPlateau(optimizer, patience=4, factor=0.25) min_loss = sys.maxsize criterion = nn.BCEWithLogitsLoss() train_data = DataLoader(TrainDataset(x_cities, mask_dir), batch_size=4, num_workers=4, shuffle=True) test_data = DataLoader(TestDataset(y_city, mask_dir), batch_size=6, num_workers=4) for epoch in range(epochs): print(f'Epoch {epoch}, lr {optimizer.param_groups[0]["lr"]}') print(f" Training") losses = [] ious = [] jaccs = [] batch_iterator = enumerate(train_data) model = model.train().cuda() for i, (x, y) in tqdm(batch_iterator): optimizer.zero_grad() x = x.cuda() y = y.cuda() y_real = y.view(-1).float() y_pred = model(x) y_pred_probs = torch.sigmoid(y_pred).view(-1) loss = 0.75 * criterion(y_pred.view( -1), y_real) + 0.25 * dice_loss(y_pred_probs, y_real) iou_ = iou(y_pred_probs.float(), y_real.byte()) jacc_ = jaccard(y_pred_probs.float(), y_real) ious.append(iou_.item()) losses.append(loss.item()) jaccs.append(jacc_.item()) loss.backward() optimizer.step() print( f"Loss: {np.mean(losses)}, IOU: {np.mean(ious)}, jacc: {np.mean(jaccs)}" ) model = model.eval().cuda() losses = [] ious = [] jaccs = [] with torch.no_grad(): batch_iterator = enumerate(test_data) for i, (x, y) in tqdm(batch_iterator): x = x.cuda() y = y.cuda() y_real = y.view(-1).float() y_pred = model(x) y_pred_probs = torch.sigmoid(y_pred).view(-1) loss = 0.75 * criterion(y_pred.view( -1), y_real) + 0.25 * dice_loss(y_pred_probs, y_real) iou_ = iou(y_pred_probs.float(), y_real.byte()) jacc_ = jaccard(y_pred_probs.float(), y_real) ious.append(iou_.item()) losses.append(loss.item()) jaccs.append(jacc_.item()) test_loss = np.mean(losses) print( f"Loss: {np.mean(losses)}, IOU: {np.mean(ious)}, jacc: {np.mean(jaccs)}" ) scheduler.step(test_loss) if test_loss < min_loss: min_loss = test_loss save_model(model, epoch, models_dir / y_city[0].name)
) partition = 'val' ultrasound_val = UltrasoundDataShapeNet(ROOT_DIR, partition) val_loader = torch.utils.data.DataLoader(ultrasound_val, batch_size=1, shuffle=False) # # Create model model = ShapeUNet((4, 544, 544)) weights = torch.load( 'weights_shape_net/Unet_manually_corrupted_lumen0.113253.pth') model.load_state_dict(weights) model.to(device) # Specify optimizer and criterion lr = 1e-4 optimizer = Adam(model.parameters(), lr=lr) NUM_OF_EPOCHS = 50 train(model, train_loader, val_loader, optimizer, NUM_OF_EPOCHS, 'weights_shape_net/') # import torch # import torch.nn as nn # from torch.optim import Adam # import torch.nn.functional as F # from data_loader import UltrasoundData # from shape_net import ShapeUNet # from tqdm import tqdm # import time
def train_experiment(device): with TemporaryDirectory() as logdir: # 1. train and valid loaders transforms = Compose([ToTensor(), Normalize((0.1307, ), (0.3081, ))]) train_dataset = datasets.MnistMLDataset(root=os.getcwd(), download=True, transform=transforms) sampler = data.BalanceBatchSampler(labels=train_dataset.get_labels(), p=5, k=10) train_loader = DataLoader(dataset=train_dataset, sampler=sampler, batch_size=sampler.batch_size) valid_dataset = datasets.MnistQGDataset(root=os.getcwd(), transform=transforms, gallery_fraq=0.2) valid_loader = DataLoader(dataset=valid_dataset, batch_size=1024) # 2. model and optimizer model = models.MnistSimpleNet(out_features=16) optimizer = Adam(model.parameters(), lr=0.001) # 3. criterion with triplets sampling sampler_inbatch = data.HardTripletsSampler(norm_required=False) criterion = nn.TripletMarginLossWithSampler( margin=0.5, sampler_inbatch=sampler_inbatch) # 4. training with catalyst Runner class CustomRunner(dl.SupervisedRunner): def handle_batch(self, batch) -> None: if self.is_train_loader: images, targets = batch["features"].float( ), batch["targets"].long() features = self.model(images) self.batch = { "embeddings": features, "targets": targets, } else: images, targets, is_query = ( batch["features"].float(), batch["targets"].long(), batch["is_query"].bool(), ) features = self.model(images) self.batch = { "embeddings": features, "targets": targets, "is_query": is_query, } callbacks = [ dl.ControlFlowCallback( dl.CriterionCallback(input_key="embeddings", target_key="targets", metric_key="loss"), loaders="train", ), dl.ControlFlowCallback( dl.CMCScoreCallback( embeddings_key="embeddings", labels_key="targets", is_query_key="is_query", topk_args=[1], ), loaders="valid", ), dl.PeriodicLoaderCallback(valid_loader_key="valid", valid_metric_key="cmc01", minimize=False, valid=2), ] runner = CustomRunner(input_key="features", output_key="embeddings") runner.train( engine=dl.DeviceEngine(device), model=model, criterion=criterion, optimizer=optimizer, callbacks=callbacks, loaders={ "train": train_loader, "valid": valid_loader }, verbose=False, logdir=logdir, valid_loader="valid", valid_metric="cmc01", minimize_valid_metric=False, num_epochs=2, )
def configure_optimizers(self): return Adam(self.parameters(), lr=self.config["lr"])
def LookaheadAdam(params, alpha=0.5, k=6, *args, **kwargs): adam = Adam(params, *args, **kwargs) return Lookahead(adam, alpha, k)
def sac(env_fn, actor_critic=core.MLPActorCritic, ac_kwargs=dict(), seed=0, steps_per_epoch=4000, epochs=100, replay_size=int(1e6), gamma=0.99, polyak=0.995, lr=1e-3, alpha=0.2, batch_size=100, start_steps=10000, update_after=1000, update_every=50, num_test_episodes=10, max_ep_len=1000, logger_kwargs=dict(), save_freq=1): """ Soft Actor-Critic (SAC) Args: env_fn : A function which creates a copy of the environment. The environment must satisfy the OpenAI Gym API. actor_critic: The constructor method for a PyTorch Module with an ``act`` method, a ``pi`` module, a ``q1`` module, and a ``q2`` module. The ``act`` method and ``pi`` module should accept batches of observations as inputs, and ``q1`` and ``q2`` should accept a batch of observations and a batch of actions as inputs. When called, ``act``, ``q1``, and ``q2`` should return: =========== ================ ====================================== Call Output Shape Description =========== ================ ====================================== ``act`` (batch, act_dim) | Numpy array of actions for each | observation. ``q1`` (batch,) | Tensor containing one current estimate | of Q* for the provided observations | and actions. (Critical: make sure to | flatten this!) ``q2`` (batch,) | Tensor containing the other current | estimate of Q* for the provided observations | and actions. (Critical: make sure to | flatten this!) =========== ================ ====================================== Calling ``pi`` should return: =========== ================ ====================================== Symbol Shape Description =========== ================ ====================================== ``a`` (batch, act_dim) | Tensor containing actions from policy | given observations. ``logp_pi`` (batch,) | Tensor containing log probabilities of | actions in ``a``. Importantly: gradients | should be able to flow back into ``a``. =========== ================ ====================================== ac_kwargs (dict): Any kwargs appropriate for the ActorCritic object you provided to SAC. seed (int): Seed for random number generators. steps_per_epoch (int): Number of steps of interaction (state-action pairs) for the agent and the environment in each epoch. epochs (int): Number of epochs to run and train agent. replay_size (int): Maximum length of replay buffer. gamma (float): Discount factor. (Always between 0 and 1.) polyak (float): Interpolation factor in polyak averaging for target networks. Target networks are updated towards main networks according to: .. math:: \\theta_{\\text{targ}} \\leftarrow \\rho \\theta_{\\text{targ}} + (1-\\rho) \\theta where :math:`\\rho` is polyak. (Always between 0 and 1, usually close to 1.) lr (float): Learning rate (used for both policy and value learning). alpha (float): Entropy regularization coefficient. (Equivalent to inverse of reward scale in the original SAC paper.) batch_size (int): Minibatch size for SGD. start_steps (int): Number of steps for uniform-random action selection, before running real policy. Helps exploration. update_after (int): Number of env interactions to collect before starting to do gradient descent updates. Ensures replay buffer is full enough for useful updates. update_every (int): Number of env interactions that should elapse between gradient descent updates. Note: Regardless of how long you wait between updates, the ratio of env steps to gradient steps is locked to 1. num_test_episodes (int): Number of episodes to test the deterministic policy at the end of each epoch. max_ep_len (int): Maximum length of trajectory / episode / rollout. logger_kwargs (dict): Keyword args for EpochLogger. save_freq (int): How often (in terms of gap between epochs) to save the current policy and value function. """ logger = EpochLogger(**logger_kwargs) logger.save_config(locals()) torch.manual_seed(seed) np.random.seed(seed) env, test_env = env_fn(), env_fn() obs_dim = env.observation_space.shape act_dim = env.action_space.shape[0] # Action limit for clamping: critically, assumes all dimensions share the same bound! act_limit = env.action_space.high[0] # Create actor-critic module and target networks ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs) ac_targ = deepcopy(ac) # Freeze target networks with respect to optimizers (only update via polyak averaging) for p in ac_targ.parameters(): p.requires_grad = False # List of parameters for both Q-networks (save this for convenience) q_params = itertools.chain(ac.q1.parameters(), ac.q2.parameters()) # Experience buffer replay_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=replay_size) # Count variables (protip: try to get a feel for how different size networks behave!) var_counts = tuple( core.count_vars(module) for module in [ac.pi, ac.q1, ac.q2]) logger.log('\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' % var_counts) # Set up function for computing SAC Q-losses def compute_loss_q(data): o, a, r, o2, d = data['obs'], data['act'], data['rew'], data[ 'obs2'], data['done'] q1 = ac.q1(o, a) q2 = ac.q2(o, a) # Bellman backup for Q functions with torch.no_grad(): # Target actions come from *current* policy a2, logp_a2 = ac.pi(o2) # Target Q-values q1_pi_targ = ac_targ.q1(o2, a2) q2_pi_targ = ac_targ.q2(o2, a2) q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ) backup = r + gamma * (1 - d) * (q_pi_targ - alpha * logp_a2) # MSE loss against Bellman backup loss_q1 = ((q1 - backup)**2).mean() loss_q2 = ((q2 - backup)**2).mean() loss_q = loss_q1 + loss_q2 # Useful info for logging q_info = dict(Q1Vals=q1.detach().numpy(), Q2Vals=q2.detach().numpy()) return loss_q, q_info # Set up function for computing SAC pi loss def compute_loss_pi(data): o = data['obs'] pi, logp_pi = ac.pi(o) q1_pi = ac.q1(o, pi) q2_pi = ac.q2(o, pi) q_pi = torch.min(q1_pi, q2_pi) # Entropy-regularized policy loss loss_pi = (alpha * logp_pi - q_pi).mean() # Useful info for logging pi_info = dict(LogPi=logp_pi.detach().numpy()) return loss_pi, pi_info # Set up optimizers for policy and q-function pi_optimizer = Adam(ac.pi.parameters(), lr=lr) q_optimizer = Adam(q_params, lr=lr) # Set up model saving logger.setup_pytorch_saver(ac) def update(data): # First run one gradient descent step for Q1 and Q2 q_optimizer.zero_grad() loss_q, q_info = compute_loss_q(data) loss_q.backward() q_optimizer.step() # Record things logger.store(LossQ=loss_q.item(), **q_info) # Freeze Q-networks so you don't waste computational effort # computing gradients for them during the policy learning step. for p in q_params: p.requires_grad = False # Next run one gradient descent step for pi. pi_optimizer.zero_grad() loss_pi, pi_info = compute_loss_pi(data) loss_pi.backward() pi_optimizer.step() # Unfreeze Q-networks so you can optimize it at next DDPG step. for p in q_params: p.requires_grad = True # Record things logger.store(LossPi=loss_pi.item(), **pi_info) # Finally, update target networks by polyak averaging. with torch.no_grad(): for p, p_targ in zip(ac.parameters(), ac_targ.parameters()): # NB: We use an in-place operations "mul_", "add_" to update target # params, as opposed to "mul" and "add", which would make new tensors. p_targ.data.mul_(polyak) p_targ.data.add_((1 - polyak) * p.data) def get_action(o, deterministic=False): return ac.act(torch.as_tensor(o, dtype=torch.float32), deterministic) def test_agent(): for j in range(num_test_episodes): o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0 while not (d or (ep_len == max_ep_len)): # Take deterministic actions at test time o, r, d, _ = test_env.step(get_action(o, True)) ep_ret += r ep_len += 1 logger.store(TestEpRet=ep_ret, TestEpLen=ep_len) # Prepare for interaction with environment total_steps = steps_per_epoch * epochs start_time = time.time() o, ep_ret, ep_len = env.reset(), 0, 0 # Main loop: collect experience in env and update/log each epoch for t in range(total_steps): # Until start_steps have elapsed, randomly sample actions # from a uniform distribution for better exploration. Afterwards, # use the learned policy. if t > start_steps: a = get_action(o) else: a = env.action_space.sample() # Step the env o2, r, d, _ = env.step(a) ep_ret += r ep_len += 1 # Ignore the "done" signal if it comes from hitting the time # horizon (that is, when it's an artificial terminal signal # that isn't based on the agent's state) d = False if ep_len == max_ep_len else d # Store experience to replay buffer replay_buffer.store(o, a, r, o2, d) # Super critical, easy to overlook step: make sure to update # most recent observation! o = o2 # End of trajectory handling if d or (ep_len == max_ep_len): logger.store(EpRet=ep_ret, EpLen=ep_len) o, ep_ret, ep_len = env.reset(), 0, 0 # Update handling if t >= update_after and t % update_every == 0: for j in range(update_every): batch = replay_buffer.sample_batch(batch_size) update(data=batch) # End of epoch handling if (t + 1) % steps_per_epoch == 0: epoch = (t + 1) // steps_per_epoch # Save model if (epoch % save_freq == 0) or (epoch == epochs): logger.save_state({'env': env}, None) # Test the performance of the deterministic version of the agent. test_agent() # Log info about epoch logger.log_tabular('Epoch', epoch) logger.log_tabular('EpRet', with_min_and_max=True) logger.log_tabular('TestEpRet', with_min_and_max=True) logger.log_tabular('EpLen', average_only=True) logger.log_tabular('TestEpLen', average_only=True) logger.log_tabular('TotalEnvInteracts', t) logger.log_tabular('Q1Vals', with_min_and_max=True) logger.log_tabular('Q2Vals', with_min_and_max=True) logger.log_tabular('LogPi', with_min_and_max=True) logger.log_tabular('LossPi', average_only=True) logger.log_tabular('LossQ', average_only=True) logger.log_tabular('Time', time.time() - start_time) logger.dump_tabular()
wordVec = Word2Vec.load('wv.h5') wordVec = wordVec.wv data_label = Data(x_label, y, wordVec) train_loader_label = DataLoader(data_label, batch_size=BATCH_SIZE, shuffle=True) print('---------------------Building Model-----------------------') #Model1 model01 = RNN01() model01.cuda() print(model01) optim = Adam(model01.parameters(), lr=LR, weight_decay=0.001) loss_function = nn.CrossEntropyLoss() loss_function.cuda() for epoch in range(EPOCH): iter_start = time.time() total_loss = 0 total_acc = 0 for iter, (x, y) in enumerate(train_loader_label): x, y = Variable(x.cuda()), Variable(y.cuda()) output = model01(x) loss = loss_function(output, y) optim.zero_grad() loss.backward() optim.step()
def td3(env_fn, actor_critic=core.MLPActorCritic, ac_kwargs=dict(), seed=0, steps_per_epoch=4000, epochs=100, replay_size=int(1e6), gamma=0.99, polyak=0.995, pi_lr=1e-3, q_lr=1e-3, batch_size=100, start_steps=10000, update_after=1000, update_every=50, act_noise=0.1, target_noise=0.2, noise_clip=0.5, policy_delay=2, num_test_episodes=10, max_ep_len=1000, logger_kwargs=dict(), save_freq=1): logger = EpochLogger(**logger_kwargs) logger.save_config(locals()) torch.manual_seed(seed) np.random.seed(seed) env, test_env = env_fn(), env_fn() obs_dim = env.observation_space.shape act_dim = env.action_space.shape[0] # 动作取值上界 act_limit = env.action_space.high[0] # 创建原网络和 target 网络 ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs) ac_targ = deepcopy(ac) # 冻结 target 网络 for p in ac_targ.parameters(): p.requires_grad = False q_params = itertools.chain(ac.q1.parameters(), ac.q2.parameters()) # 经验池 replay_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=replay_size) var_counts = tuple( core.count_vars(module) for module in [ac.pi, ac.q1, ac.q2]) logger.log('\nNumber of parameters: \t pi: %d, \t q1: %d, \t q2: %d\n' % var_counts) # TD3 Q loss 计算 def compute_loss_q(data): o, a, r, o2, d = data['obs'], data['act'], data['rew'], data[ 'obs2'], data['done'] # 俩原网络的 Q 值 q1 = ac.q1(o, a) q2 = ac.q2(o, a) with torch.no_grad(): pi_targ = ac_targ.pi(o2) # 目标 ACTOR 网络平滑 epsilon = torch.randn_like(pi_targ) * target_noise epsilon = torch.clamp(epsilon, -noise_clip, noise_clip) a2 = pi_targ + epsilon a2 = torch.clamp(a2, -act_limit, act_limit) # 目标 Q 值计算 q1_pi_targ = ac_targ.q1(o2, a2) q2_pi_targ = ac_targ.q2(o2, a2) q_pi_targ = torch.min(q1_pi_targ, q2_pi_targ) backup = r + gamma * (1 - d) * q_pi_targ # MSE 损失函数 loss_q1 = ((q1 - backup)**2).mean() loss_q2 = ((q2 - backup)**2).mean() loss_q = loss_q1 + loss_q2 # logging 用的 loss_info = dict(Q1Vals=q1.detach().numpy(), Q2Vals=q2.detach().numpy()) return loss_q, loss_info # actor 部分 loss def compute_loss_pi(data): o = data['obs'] q1_pi = ac.q1(o, ac.pi(o)) return -q1_pi.mean() # 优化器 pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr) q_optimizer = Adam(q_params, lr=q_lr) logger.setup_pytorch_saver(ac) def update(data, timer): # 对 Q1 Q2 梯度下降 q_optimizer.zero_grad() loss_q, loss_info = compute_loss_q(data) loss_q.backward() q_optimizer.step() # 记录 logger.store(LossQ=loss_q.item(), **loss_info) # actor 更新 if timer % policy_delay == 0: for p in q_params: p.requires_grad = False pi_optimizer.zero_grad() loss_pi = compute_loss_pi(data) loss_pi.backward() pi_optimizer.step() for p in q_params: p.requires_grad = True # 记录 logger.store(LossPi=loss_pi.item()) with torch.no_grad(): for p, p_targ in zip(ac.parameters(), ac_targ.parameters()): p_targ.data.mul_(polyak) p_targ.data.add_((1 - polyak) * p.data) def get_action(o, noise_scale): a = ac.act(torch.as_tensor(o, dtype=torch.float32)) a += noise_scale * np.random.randn(act_dim) return np.clip(a, -act_limit, act_limit) def test_agent(): for j in range(num_test_episodes): o, d, ep_ret, ep_len = test_env.reset(), False, 0, 0 while not (d or (ep_len == max_ep_len)): # Take deterministic actions at test time (noise_scale=0) o, r, d, _ = test_env.step(get_action(o, 0)) ep_ret += r ep_len += 1 logger.store(TestEpRet=ep_ret, TestEpLen=ep_len) total_steps = steps_per_epoch * epochs start_time = time.time() o, ep_ret, ep_len = env.reset(), 0, 0 for t in range(total_steps): if t > start_steps: a = get_action(o, act_noise) else: a = env.action_space.sample() o2, r, d, _ = env.step(a) ep_ret += r ep_len += 1 d = False if ep_len == max_ep_len else d replay_buffer.store(o, a, r, o2, d) o = o2 if d or (ep_len == max_ep_len): logger.store(EpRet=ep_ret, EpLen=ep_len) o, ep_ret, ep_len = env.reset(), 0, 0 if t >= update_after and t % update_every == 0: for j in range(update_every): batch = replay_buffer.sample_batch(batch_size) update(data=batch, timer=j) if (t + 1) % steps_per_epoch == 0: epoch = (t + 1) // steps_per_epoch if (epoch % save_freq == 0) or (epoch == epochs): logger.save_state({'env': env}, None) test_agent() logger.log_tabular('Epoch', epoch) logger.log_tabular('EpRet', with_min_and_max=True) logger.log_tabular('TestEpRet', with_min_and_max=True) logger.log_tabular('EpLen', average_only=True) logger.log_tabular('TestEpLen', average_only=True) logger.log_tabular('TotalEnvInteracts', t) logger.log_tabular('Q1Vals', with_min_and_max=True) logger.log_tabular('Q2Vals', with_min_and_max=True) logger.log_tabular('LossPi', average_only=True) logger.log_tabular('LossQ', average_only=True) logger.log_tabular('Time', time.time() - start_time) logger.dump_tabular()
def train(visualize=True): torch.manual_seed(config['seed']) torch.cuda.manual_seed(config['seed']) random.seed(config['seed']) device = ("cuda" if torch.cuda.is_available() else "cpu") generator = Generator().to(device) restorer = Restorer().to(device) params = list(generator.parameters()) + list(restorer.parameters()) # joint optimization optimizer = Adam(params, lr=config['lr'], weight_decay=config['weight_decay']) # Start from previous session if config['checkpoint_path'] is not None: checkpoint = torch.load(config['checkpoint_path']) generator.load_state_dict(checkpoint["generator"]) restorer.load_state_dict(checkpoint["restorer"]) optimizer.load_state_dict(checkpoint["optimize"]) generator.train() restorer.train() # loss parameters alpha = config['loss']['alpha'] beta = config['loss']['beta'] gamma = config['loss']['gamma'] l1_loss = L1Loss().to(device) # Dataset dataset = CustomDataset() dataloader = DataLoader(dataset, batch_size=config['batch_size'], shuffle=True) # Training Loop print("Training begins") logger = Logger(visualize=visualize) # logs and visdom plots, images batch_count, checkpoint_count = 0, 0 for epoch in range(config['n_epochs']): print("Starting Epoch #{}...".format(epoch)) for batch_i, images in enumerate(dataloader): torch.cuda.empty_cache() optimizer.zero_grad() photos = images[0].to(device) true_sketches = images[1].to(device) deteriorated_sketches = images[2].to(device) generated_sketches = generator(photos) corrected_sketches = restorer(generated_sketches) corrected_deteriorated = restorer(deteriorated_sketches) loss_base = loss_w1(generated_sketches, true_sketches, gamma) # L_base loss_aux = loss_w1(corrected_sketches, true_sketches, gamma) # L_aux loss_res = l1_loss(corrected_deteriorated, true_sketches) # L_res loss_joint = loss_base + alpha * loss_aux + beta * loss_res # L_joint loss_joint.backward() optimizer.step() # logs batch losses logger.log_iteration(epoch, batch_i, loss_base.item(), loss_res.item(), loss_aux.item(), loss_joint.item() ) if visualize and batch_count % config['sample_interval'] == 0: logger.draw( corrected_sketches[0].data.numpy() * 255, true_sketches[0].data.numpy() * 255, corrected_deteriorated[0].data.numpy() * 255, dataset.unnormalize(photos[0]).data.numpy() ) # draws a sample on Visdom # checkpoint save if batch_count % config['save_checkpoint_interval'] == 0: torch.save({ "generator": generator.state_dict(), "restorer": restorer.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch }, "{}/checkpoint_{}.pth".format(config['checkpoint_dir'], checkpoint_count)) checkpoint_count += 1 batch_count += 1 if visualize: logger.plot_epoch(epoch) # plots the average losses on Visdom
def __init__(self, args): # ------------------------- Initialization --------------------------- # region Global Settings os.environ['CUDA_VISIBLE_DEVICES'] = args.device_ids.strip() self.device_count = len(args.device_ids.strip().split(",")) output_date = time.strftime("%Y%m%d%H%M%S") self.output_dir = os.path.join(args.output_dir, output_date) if not os.path.exists(self.output_dir): os.mkdir(self.output_dir) setlogger(os.path.join(self.output_dir, "train.log")) # endregion # region Vocabulary Settings self.words2index, self.index2words = load_vocabulary( args.voc_path) # Vocabulary dict map self.vocab_len = len(self.words2index.keys()) self.start_symbol = self.words2index["<START>"] self.pad_symbol = self.words2index["<PAD>"] self.end_symbol = self.words2index["<END>"] # endregion # region Model Settings model_nonedp = TransformerCNN(trg_vocab=self.vocab_len).to( device) # none data parallel logging.info(model_nonedp) d_model = model_nonedp.d_model if args.resume != '': model_nonedp.load_state_dict(torch.load(args.resume)) if self.device_count > 1: self.model = nn.DataParallel(model_nonedp) # Not Supported else: self.model = model_nonedp # endregion # region Data Settings self.datasets = { split: Dataset(args.anno_path, words2index=self.words2index, image_root=args.image_root, split=split) for split in ['train', 'val', 'test'] } self.samplers = { split: Sampler( batch_size=args.batch_size, dataset=self.datasets[split], ) for split in ['train', 'val', 'test'] } self.dataloaders = { split: DataLoader( self.datasets[split], # batch_size = args.batch_size, shuffle = True if split=='train' else False, batch_sampler=self.samplers[split], pin_memory=True, num_workers=4, ) for split in ['train', 'val', 'test'] } db_maxlen = max([ self.datasets[split].maxlength for split in ['train', 'val', 'test'] ]) # endregion # region Optimizer Settings if args.opt_type == "adam": from torch.optim import Adam self.optimizer = OptWrapper( d_model, args.noamopt_factor, args.noamopt_warmup, Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=args.lr, betas=(0.9, 0.98), eps=1e-9, weight_decay=args.weight_decay)) elif args.opt_type == "adamw": self.optimizer = OptWrapper( d_model, args.noamopt_factor, args.noamopt_warmup, AdamW(self.model.parameters(), lr=args.lr, betas=(0.9, 0.98), eps=1e-9, weight_decay=args.weight_decay)) elif args.opt_type == "sgd": from torch.optim import SGD self.optimizer = OptWrapper( d_model, args.noamopt_factor, args.noamopt_warmup, SGD(self.model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)) elif args.opt_type == "sgdw": global optimizer self.optimizer = OptWrapper( d_model, args.noamopt_factor, args.noamopt_warmup, SGDW(self.model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)) elif args.opt_type == "adagrad": from torch.optim import Adagrad self.optimizer = OptWrapper( d_model, args.noamopt_factor, args.noamopt_warmup, Adagrad(self.model.parameters(), lr=args.lr, eps=1e-9, weight_decay=args.weight_decay)) else: ### Add other custom optimizers here.. raise ValueError(f"Invalid opt_type:{args.opt_type}") # endregion # region Criterion(Loss) Settings if args.crit_type == "crossentropy": self.criterion = CrossEntropyLoss(generator=model_nonedp.generator, ignore_index=self.pad_symbol) elif args.crit_type == "kldivloss": self.criterion = LabelSmoothingLoss( generator=model_nonedp.generator, ignore_index=self.pad_symbol) else: ### Add other custom criterions here.. raise ValueError(f"Invalid crit_type:{args.crit_type}") # endregion # region Translator Settings if args.val_translator == "greedy": self.translator = GreedyTranslator( dataparallel=True if self.device_count > 1 else False, max_len=db_maxlen, start_symbol=self.start_symbol) else: ### Add other custom translators here.. raise ValueError(f"Invalid val_translator:{args.val_translator}") # endregion # region Validate Metrics Settings if args.val_metrics == "bleu": self.metricwapper = MetricWrapper(self.index2words, self.start_symbol, self.end_symbol, self.pad_symbol, metric=Bleu()) else: ### Add other custom metrics here.. raise ValueError(f"Invalid val_metrics:{args.val_metrics}") # endregion # ---------------------------- Training ----------------------------- # Pre-training: Design custom tasks in pre-training phase if args.pretrain: self.samplers['train'].batch_size = args.pretrain_batch_size self.pretrain(args.pretrain_max_epoch, args.display_interval) # Training self.samplers['train'].batch_size = args.batch_size self.train(args.max_epoch, args.val_interval, args.display_interval)
x = F.max_pool2d(F.relu(self.conv4(x)), 2) x = x.view(-1, 64 * 30 * 40) x = F.relu(self.fc1(x)) x = self.fc2(x) return x net = Net() params = list(net.parameters()) params[1] criterion = nn.MSELoss() optim = Adam(net.parameters(), lr=0.001) for epoch in range(1): for i in range(len(train_data.csv) // train_data.batch_size): x, y = train_data.get_next_batch() x, y = Variable(torch.from_numpy(x)), Variable(torch.from_numpy(y)) optim.zero_grad() out = net(x) loss = criterion(out, y) loss.backward() optim.step() if (i % 100 == 0): print("Epoch :{} Iter : {} Loss : {} ".format(epoch, i, loss))
help='load pretrained model (default: False)') args = parser.parse_args() batch_loader = BatchLoader() parameters = Parameters(batch_loader.vocab_size) vae = VAE(parameters.vocab_size, parameters.embed_size, parameters.latent_size, parameters.decoder_rnn_size, parameters.decoder_rnn_num_layers) if args.use_trained: vae.load_state_dict(t.load('trained_VAE')) if args.use_cuda: vae = vae.cuda() optimizer = Adam(vae.parameters(), args.learning_rate) for iteration in range(args.num_iterations): '''Train step''' input, decoder_input, target = batch_loader.next_batch( args.batch_size, 'train', args.use_cuda) target = target.view(-1) logits, aux_logits, kld = vae(args.dropout, input, decoder_input) logits = logits.view(-1, batch_loader.vocab_size) cross_entropy = F.cross_entropy(logits, target, size_average=False) aux_logits = aux_logits.view(-1, batch_loader.vocab_size) aux_cross_entropy = F.cross_entropy(aux_logits, target,
def train(args, model, enc=False): best_acc = 0 #TODO: calculate weights by processing dataset histogram (now its being set by hand from the torch values) #create a loder to run all images and calculate histogram of labels, then create weight array using class balancing weight = torch.ones(NUM_CLASSES) if (enc): weight[0] = 2.3653597831726 weight[1] = 4.4237880706787 weight[2] = 2.9691488742828 weight[3] = 5.3442072868347 weight[4] = 5.2983593940735 weight[5] = 5.2275490760803 weight[6] = 5.4394111633301 weight[7] = 5.3659925460815 weight[8] = 3.4170460700989 weight[9] = 5.2414722442627 weight[10] = 4.7376127243042 weight[11] = 5.2286224365234 weight[12] = 5.455126285553 weight[13] = 4.3019247055054 weight[14] = 5.4264230728149 weight[15] = 5.4331531524658 weight[16] = 5.433765411377 weight[17] = 5.4631009101868 weight[18] = 5.3947434425354 else: weight[0] = 2.8149201869965 weight[1] = 6.9850029945374 weight[2] = 3.7890393733978 weight[3] = 9.9428062438965 weight[4] = 9.7702074050903 weight[5] = 9.5110931396484 weight[6] = 10.311357498169 weight[7] = 10.026463508606 weight[8] = 4.6323022842407 weight[9] = 9.5608062744141 weight[10] = 7.8698215484619 weight[11] = 9.5168733596802 weight[12] = 10.373730659485 weight[13] = 6.6616044044495 weight[14] = 10.260489463806 weight[15] = 10.287888526917 weight[16] = 10.289801597595 weight[17] = 10.405355453491 weight[18] = 10.138095855713 weight[19] = 0 assert os.path.exists( args.datadir), "Error: datadir (dataset directory) could not be loaded" """ co_transform = MyCoTransform(enc, augment=True, height=args.height)#1024) co_transform_val = MyCoTransform(enc, augment=False, height=args.height)#1024) dataset_train = cityscapes(args.datadir, co_transform, 'train') dataset_val = cityscapes(args.datadir, co_transform_val, 'val') loader = DataLoader(dataset_train, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True) loader_val = DataLoader(dataset_val, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False) """ ### RELLIS-3D Dataloader loader, loader_val = custom_datasets.setup_loaders(args, enc) if args.cuda: weight = weight.cuda() criterion = CrossEntropyLoss2d(weight) print(type(criterion)) savedir = f'../save/{args.savedir}' if (enc): automated_log_path = savedir + "/automated_log_encoder.txt" modeltxtpath = savedir + "/model_encoder.txt" else: automated_log_path = savedir + "/automated_log.txt" modeltxtpath = savedir + "/model.txt" if (not os.path.exists(automated_log_path) ): #dont add first line if it exists with open(automated_log_path, "a") as myfile: myfile.write( "Epoch\t\tTrain-loss\t\tTest-loss\t\tTrain-IoU\t\tTest-IoU\t\tlearningRate" ) with open(modeltxtpath, "w") as myfile: myfile.write(str(model)) #TODO: reduce memory in first gpu: https://discuss.pytorch.org/t/multi-gpu-training-memory-usage-in-balance/4163/4 #https://github.com/pytorch/pytorch/issues/1893 #optimizer = Adam(model.parameters(), 5e-4, (0.9, 0.999), eps=1e-08, weight_decay=2e-4) ## scheduler 1 optimizer = Adam(model.parameters(), 5e-4, (0.9, 0.999), eps=1e-08, weight_decay=1e-4) ## scheduler 2 start_epoch = 1 if args.resume: #Must load weights, optimizer, epoch and best value. if enc: filenameCheckpoint = savedir + '/checkpoint_enc.pth.tar' else: filenameCheckpoint = savedir + '/checkpoint.pth.tar' assert os.path.exists( filenameCheckpoint ), "Error: resume option was used but checkpoint was not found in folder" checkpoint = torch.load(filenameCheckpoint) start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) best_acc = checkpoint['best_acc'] print("=> Loaded checkpoint at epoch {})".format(checkpoint['epoch'])) #scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5) # set up scheduler ## scheduler 1 lambda1 = lambda epoch: pow( (1 - ((epoch - 1) / args.num_epochs)), 0.9) ## scheduler 2 scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1) ## scheduler 2 if args.visualize and args.steps_plot > 0: board = Dashboard(args.port) for epoch in range(start_epoch, args.num_epochs + 1): print("----- TRAINING - EPOCH", epoch, "-----") scheduler.step(epoch) ## scheduler 2 epoch_loss = [] time_train = [] doIouTrain = args.iouTrain doIouVal = args.iouVal if (doIouTrain): iouEvalTrain = iouEval(NUM_CLASSES) usedLr = 0 for param_group in optimizer.param_groups: print("LEARNING RATE: ", param_group['lr']) usedLr = float(param_group['lr']) model.train() for step, (images, labels) in enumerate(loader): start_time = time.time() #print (labels.size()) #print (np.unique(labels.numpy())) #print("labels: ", np.unique(labels[0].numpy())) #labels = torch.ones(4, 1, 512, 1024).long() if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable(images) #print("inputs.shape = {}".format(inputs.shape)) targets = Variable(labels) #print("targets.shape = {}".format(targets.shape)) outputs = model(inputs, only_encode=enc) #print("outputs.shape = {}".format(outputs.shape)) #print("targets[:, 0].long() = {}".format(targets[:, 0].shape)) #print("targets", np.unique(targets[:, 0].cpu().data.numpy())) optimizer.zero_grad() loss = criterion(outputs, targets[:, 0]) #loss = criterion(outputs, targets[:, 0].long()) #print("loss= {}".format(loss)) loss.backward() optimizer.step() #epoch_loss.append(loss.data[0]) epoch_loss.append(loss.item()) time_train.append(time.time() - start_time) if (doIouTrain): #start_time_iou = time.time() iouEvalTrain.addBatch( outputs.max(1)[1].unsqueeze(1).data, targets.data) #print ("Time to add confusion matrix: ", time.time() - start_time_iou) #print(outputs.size()) if args.visualize and args.steps_plot > 0 and step % args.steps_plot == 0: start_time_plot = time.time() image = inputs[0].cpu().data #image[0] = image[0] * .229 + .485 #image[1] = image[1] * .224 + .456 #image[2] = image[2] * .225 + .406 #print("output", np.unique(outputs[0].cpu().max(0)[1].data.numpy())) board.image(image, f'input (epoch: {epoch}, step: {step})') if isinstance(outputs, list): #merge gpu tensors board.image( color_transform( outputs[0][0].cpu().max(0)[1].data.unsqueeze(0)), f'output (epoch: {epoch}, step: {step})') else: board.image( color_transform( outputs[0].cpu().max(0)[1].data.unsqueeze(0)), f'output (epoch: {epoch}, step: {step})') board.image(color_transform(targets[0].cpu().data), f'target (epoch: {epoch}, step: {step})') print("Time to paint images: ", time.time() - start_time_plot) if args.steps_loss > 0 and step % args.steps_loss == 0: average = sum(epoch_loss) / len(epoch_loss) print( f'loss: {average:0.4} (epoch: {epoch}, step: {step})', "// Avg time/img: %.4f s" % (sum(time_train) / len(time_train) / args.batch_size)) average_epoch_loss_train = sum(epoch_loss) / len(epoch_loss) iouTrain = 0 if (doIouTrain): iouTrain, iou_classes = iouEvalTrain.getIoU() iouStr = getColorEntry(iouTrain) + '{:0.2f}'.format( iouTrain * 100) + '\033[0m' print("EPOCH IoU on TRAIN set: ", iouStr, "%") #Validate on 500 val images after each epoch of training print("----- VALIDATING - EPOCH", epoch, "-----") model.eval() epoch_loss_val = [] time_val = [] if (doIouVal): iouEvalVal = iouEval(NUM_CLASSES) for step, (images, labels) in enumerate(loader_val): start_time = time.time() if args.cuda: images = images.cuda() labels = labels.cuda() inputs = Variable( images, volatile=True ) #volatile flag makes it free backward or outputs for eval targets = Variable(labels, volatile=True) outputs = model(inputs, only_encode=enc) loss = criterion(outputs, targets[:, 0]) #epoch_loss_val.append(loss.data[0]) epoch_loss_val.append(loss.item()) time_val.append(time.time() - start_time) #Add batch to calculate TP, FP and FN for iou estimation if (doIouVal): #start_time_iou = time.time() iouEvalVal.addBatch( outputs.max(1)[1].unsqueeze(1).data, targets.data) #print ("Time to add confusion matrix: ", time.time() - start_time_iou) if args.visualize and args.steps_plot > 0 and step % args.steps_plot == 0: start_time_plot = time.time() image = inputs[0].cpu().data board.image(image, f'VAL input (epoch: {epoch}, step: {step})') if isinstance(outputs, list): #merge gpu tensors board.image( color_transform( outputs[0][0].cpu().max(0)[1].data.unsqueeze(0)), f'VAL output (epoch: {epoch}, step: {step})') else: board.image( color_transform( outputs[0].cpu().max(0)[1].data.unsqueeze(0)), f'VAL output (epoch: {epoch}, step: {step})') board.image(color_transform(targets[0].cpu().data), f'VAL target (epoch: {epoch}, step: {step})') print("Time to paint images: ", time.time() - start_time_plot) if args.steps_loss > 0 and step % args.steps_loss == 0: average = sum(epoch_loss_val) / len(epoch_loss_val) print( f'VAL loss: {average:0.4} (epoch: {epoch}, step: {step})', "// Avg time/img: %.4f s" % (sum(time_val) / len(time_val) / args.batch_size)) average_epoch_loss_val = sum(epoch_loss_val) / len(epoch_loss_val) #scheduler.step(average_epoch_loss_val, epoch) ## scheduler 1 # update lr if needed iouVal = 0 if (doIouVal): iouVal, iou_classes = iouEvalVal.getIoU() iouStr = getColorEntry(iouVal) + '{:0.2f}'.format( iouVal * 100) + '\033[0m' print("EPOCH IoU on VAL set: ", iouStr, "%") # remember best valIoU and save checkpoint if iouVal == 0: current_acc = -average_epoch_loss_val else: current_acc = iouVal is_best = current_acc > best_acc best_acc = max(current_acc, best_acc) if enc: filenameCheckpoint = savedir + '/checkpoint_enc.pth.tar' filenameBest = savedir + '/model_best_enc.pth.tar' else: filenameCheckpoint = savedir + '/checkpoint.pth.tar' filenameBest = savedir + '/model_best.pth.tar' save_checkpoint( { 'epoch': epoch + 1, 'arch': str(model), 'state_dict': model.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, is_best, filenameCheckpoint, filenameBest) #SAVE MODEL AFTER EPOCH if (enc): filename = f'{savedir}/model_encoder-{epoch:03}.pth' filenamebest = f'{savedir}/model_encoder_best.pth' else: filename = f'{savedir}/model-{epoch:03}.pth' filenamebest = f'{savedir}/model_best.pth' if args.epochs_save > 0 and step > 0 and step % args.epochs_save == 0: torch.save(model.state_dict(), filename) print(f'save: {filename} (epoch: {epoch})') if (is_best): torch.save(model.state_dict(), filenamebest) print(f'save: {filenamebest} (epoch: {epoch})') if (not enc): with open(savedir + "/best.txt", "w") as myfile: myfile.write("Best epoch is %d, with Val-IoU= %.4f" % (epoch, iouVal)) else: with open(savedir + "/best_encoder.txt", "w") as myfile: myfile.write("Best epoch is %d, with Val-IoU= %.4f" % (epoch, iouVal)) #SAVE TO FILE A ROW WITH THE EPOCH RESULT (train loss, val loss, train IoU, val IoU) #Epoch Train-loss Test-loss Train-IoU Test-IoU learningRate with open(automated_log_path, "a") as myfile: myfile.write("\n%d\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.8f" % (epoch, average_epoch_loss_train, average_epoch_loss_val, iouTrain, iouVal, usedLr)) return (model) #return model (convenience for encoder-decoder training)
def train(bundles, model1, device, mode, model2, batch_size, num_epoch, gradient_accumulation_steps, lr1, lr2, alpha): '''Train Sys1 and Sys2 models. Train models by task #1(tensors) and task #2(bundle). Args: bundles (list): List of bundles. model1 (BertForMultiHopQuestionAnswering): System 1 model. device (torch.device): The device which models and data are on. mode (str): Defaults to 'tensors'. Task identifier('tensors' or 'bundle'). model2 (CognitiveGNN): System 2 model. batch_size (int): Defaults to 4. num_epoch (int): Defaults to 1. gradient_accumulation_steps (int): Defaults to 1. lr1 (float): Defaults to 1e-4. Learning rate for Sys1. lr2 (float): Defaults to 1e-4. Learning rate for Sys2. alpha (float): Defaults to 0.2. Balance factor for loss of two systems. Returns: ([type], [type]): Trained models. ''' # Prepare optimizer for Sys1 param_optimizer = list(model1.named_parameters()) # hack to remove pooler, which is not used. param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]] no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight'] optimizer_grouped_parameters = [{ 'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01 }, { 'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0 }] num_batch, dataloader = homebrew_data_loader(bundles, mode=mode, batch_size=batch_size) num_steps = num_batch * num_epoch global_step = 0 opt1 = BertAdam(optimizer_grouped_parameters, lr=lr1, warmup=0.1, t_total=num_steps) model1.to(device) model1.train() # Prepare optimizer for Sys2 if mode == 'bundle': opt2 = Adam(model2.parameters(), lr=lr2) model2.to(device) model2.train() warmed = False # warmup for jointly training for epoch in trange(num_epoch, desc='Epoch'): ans_mean, hop_mean = WindowMean(), WindowMean() opt1.zero_grad() if mode == 'bundle': final_mean = WindowMean() opt2.zero_grad() tqdm_obj = tqdm(dataloader, total=num_batch) for step, batch in enumerate(tqdm_obj): try: if mode == 'tensors': batch = tuple(t.to(device) for t in batch) hop_loss, ans_loss, pooled_output = model1(*batch) hop_loss, ans_loss = hop_loss.mean(), ans_loss.mean() pooled_output.detach() loss = ans_loss + hop_loss elif mode == 'bundle': hop_loss, ans_loss, final_loss = model2( batch, model1, device) hop_loss, ans_loss = hop_loss.mean(), ans_loss.mean() loss = ans_loss + hop_loss + alpha * final_loss loss.backward() if (step + 1) % gradient_accumulation_steps == 0: # modify learning rate with special warm up BERT uses. From BERT pytorch examples lr_this_step = lr1 * warmup_linear(global_step / num_steps, warmup=0.1) for param_group in opt1.param_groups: param_group['lr'] = lr_this_step global_step += 1 if mode == 'bundle': opt2.step() opt2.zero_grad() final_mean_loss = final_mean.update(final_loss.item()) tqdm_obj.set_description( 'ans_loss: {:.2f}, hop_loss: {:.2f}, final_loss: {:.2f}' .format(ans_mean.update(ans_loss.item()), hop_mean.update(hop_loss.item()), final_mean_loss)) # During warming period, model1 is frozen and model2 is trained to normal weights if final_mean_loss < 0.9 and step > 100: # ugly manual hyperparam warmed = True if warmed: opt1.step() opt1.zero_grad() else: opt1.step() opt1.zero_grad() tqdm_obj.set_description( 'ans_loss: {:.2f}, hop_loss: {:.2f}'.format( ans_mean.update(ans_loss.item()), hop_mean.update(hop_loss.item()))) if step % 1000 == 0: output_model_file = './models/bert-base-uncased.bin.tmp' saved_dict = {'params1': model1.module.state_dict()} saved_dict['params2'] = model2.state_dict() torch.save(saved_dict, output_model_file) except Exception as err: traceback.print_exc() if mode == 'bundle': print(batch._id) return (model1, model2)
def train(args): device = torch.device("cuda" if args.cuda else "cpu") np.random.seed(args.seed) torch.manual_seed(args.seed) transform = transforms.Compose([ transforms.Resize(args.image_size), transforms.CenterCrop(args.image_size), transforms.ToTensor(), transforms.Lambda(lambda x: x.mul(255)) ]) train_dataset = datasets.ImageFolder(args.dataset, transform) train_loader = DataLoader(train_dataset, batch_size=args.batch_size) transformer = TransformerNet().to(device) optimizer = Adam(transformer.parameters(), args.lr) mse_loss = torch.nn.MSELoss() vgg = Vgg16(requires_grad=False).to(device) style_transform = transforms.Compose([ transforms.ToTensor(), transforms.Lambda(lambda x: x.mul(255)) ]) style = utils.load_image(args.style_image, size=args.style_size) style = style_transform(style) style = style.repeat(args.batch_size, 1, 1, 1).to(device) features_style = vgg(utils.normalize_batch(style)) gram_style = [utils.gram_matrix(y) for y in features_style] for e in range(args.epochs): transformer.train() agg_content_loss = 0. agg_style_loss = 0. count = 0 for batch_id, (x, _) in enumerate(train_loader): n_batch = len(x) count += n_batch optimizer.zero_grad() x = x.to(device) y = transformer(x) y = utils.normalize_batch(y) x = utils.normalize_batch(x) features_y = vgg(y) features_x = vgg(x) content_loss = args.content_weight * mse_loss(features_y.relu2_2, features_x.relu2_2) style_loss = 0. for ft_y, gm_s in zip(features_y, gram_style): gm_y = utils.gram_matrix(ft_y) style_loss += mse_loss(gm_y, gm_s[:n_batch, :, :]) style_loss *= args.style_weight total_loss = content_loss + style_loss total_loss.backward() optimizer.step() agg_content_loss += content_loss.item() agg_style_loss += style_loss.item() if (batch_id + 1) % args.log_interval == 0: mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.6f}\tstyle: {:.6f}\ttotal: {:.6f}".format( time.ctime(), e + 1, count, len(train_dataset), agg_content_loss / (batch_id + 1), agg_style_loss / (batch_id + 1), (agg_content_loss + agg_style_loss) / (batch_id + 1) ) print(mesg) if args.checkpoint_model_dir is not None and (batch_id + 1) % args.checkpoint_interval == 0: transformer.eval().cpu() ckpt_model_filename = "ckpt_epoch_" + str(e) + "_batch_id_" + str(batch_id + 1) + ".pth" ckpt_model_path = os.path.join(args.checkpoint_model_dir, ckpt_model_filename) torch.save(transformer.state_dict(), ckpt_model_path) transformer.to(device).train() # save model transformer.eval().cpu() save_model_filename = "epoch_" + str(args.epochs) + "_" + str(time.ctime()).replace(' ', '_') + "_" + str( args.content_weight) + "_" + str(args.style_weight) + ".model" save_model_path = os.path.join(args.save_model_dir, save_model_filename) torch.save(transformer.state_dict(), save_model_path) print("\nDone, trained model saved at", save_model_path)
def main(): # Argument parsing: parser = argparse.ArgumentParser() parser.add_argument("csvfile_dir", help="CSV file contaning the file_dir,genre data", type=str) parser.add_argument("basedir", help="Base directory for the audios", type=str) parser.add_argument("-b", "--batch_size", default=32, type=int, help="Batch size used for the training. Default 32") parser.add_argument("-e", "--epochs", default=10, type=int, help="Number of epochs to train. Default 10") parser.add_argument( "-m", "--model_dictionary_state", default="model.pt", type=str, help= "Pytorch file where to store model dictionary state. Default model.pt") args = parser.parse_args() csvfile_dir = args.csvfile_dir bsegmentdir = args.basedir if os.path.isfile(args.model_dictionary_state): response = input( "File {} already exists. Want to override it? [y]/n: ".format( args.model_dictionary_state)) if response == '' or response.lower() == 'y': m_state_dic = args.csvfile_dir else: m_state_dic = input("Type the file for model dictionary output") else: m_state_dic = args.csvfile_dir # Load the datasets daset = LatinSegments(csvfile_dir, bsegmentdir, transform=True) train_size = int(0.6 * len(daset)) tt_aux = len(daset) - train_size test_size = int(0.5 * tt_aux) valid_size = tt_aux - test_size train_dataset_d, test_dataset_d, validation_dataset_d = random_split( daset, [train_size, test_size, valid_size]) bs = args.batch_size # Batch size train_dl = DataLoader(train_dataset_d, batch_size=bs) test_dl = DataLoader(test_dataset_d, batch_size=bs) val_dl = DataLoader(validation_dataset_d, batch_size=1) model = GenreClassifier() device = t.device('cuda' if t.cuda.is_available() else 'cpu') print('Using device:', device) print('---------------') model.to(device) # Additional Info when using cuda if device.type == 'cuda': print(t.cuda.get_device_name(0)) print('Memory Usage:') print('Allocated:', round(t.cuda.memory_allocated(0) / 1024**3, 1), 'GB') print('Cached: ', round(t.cuda.memory_cached(0) / 1024**3, 1), 'GB') # The encodings genredict = { 'reggaeton': 0, 'bachata': 1, 'salsa': 2, 'merengue': 3, 'chachacha': 4 } def str_t_code(strings, diction): return t.tensor([diction.get(string) for string in strings]).to(device) loss_fun = nn.NLLLoss() # The loss function # Main loop loss_train = [] loss_test = [] cmatrixes = [] epochs = args.epochs try: for epoch in range(epochs): model.train() optimizer = Adam(model.parameters(), 1e-4 / (epoch / 4 + 1)) # Train for this epoch aux = [] print("*", "_" * 48, "*") print("| {:^14}| | {:^8} | | {:^7} |".format( 'loss', 'mini-batch(%total)', 'epoch')) for idx, sample in enumerate(train_dl): optimizer.zero_grad() spectrogram = sample[0] gcode = str_t_code(sample[1], genredict) predicted = model(spectrogram.to(device)) loss = loss_fun(predicted, gcode) loss.backward() optimizer.step() if (idx + 1) % 10 == 0: print("| {:^10.4f} | {:>5}({:4f}%) | {:^4d} |". format(loss.item(), (idx + 1), (idx + 1) / len(train_dl) * 100, epoch)) aux.append(loss.item()) loss_train.append(np.mean(aux)) # Epoch test check model.eval() with t.no_grad(): aux = [] for idx, sample in enumerate(test_dl): optimizer.zero_grad() spectrogram = sample[0] gcode = str_t_code(sample[1], genredict) predicted = model(spectrogram.to(device)) loss = loss_fun(predicted, gcode) aux.append(loss.item()) loss_test.append(np.mean(aux)) print( "######### FINISHED EPOCH {:^10d} ##########".format( epoch)) print("######## EPOCH TEST LOSS {:^10.4f} #########". format(loss_test[-1])) try: if loss_test[-2] > loss_test[-1]: t.save(model.state_dict(), m_state_dic) except IndexError: t.save(model.state_dict(), m_state_dic) #Epoch accuracy val p = [] act = [] inverse_dict = {val: key for key, val in genredict.items()} model.eval() with t.no_grad(): for sample in val_dl: predicted = model(sample[0].to(device)) predicted_g = inverse_dict[predicted.argmax(1).item()] p.append(predicted_g) act.append(sample[1][0]) confusion_mat = confusion_matrix(act, p) ok_classifications = np.sum( [confusion_mat[i][i] for i in range(confusion_mat.shape[0])]) tot_class = np.sum(confusion_mat) print("######## EPOCH ACCURACY {:^10.4f} #########". format(ok_classifications / tot_class)) cmatrixes.append(confusion_mat) np.save('loss_train.npy', loss_train) np.save('loss_test.npy', loss_test) np.save('c_matrix.npy', cmatrixes) print("Finished training!") except KeyboardInterrupt: print("Interrupting training") t.save(model.state_dict(), m_state_dic + '_interrupted')
def main(args): #Randomly initialize actor and critic networks actor = nn.Actor(input_dim=NUM_LOCATIONS, hidden_dim=ACTOR_HIDDEN, output_dim=ACTOR_OUTPUT) actor_optim = Adam(actor.parameters(), lr=ACTOR_LR) critic = nn.Critic(state_dim=NUM_LOCATIONS, action_dim=ACTOR_OUTPUT, hidden=CRTIC_HIDDEN, output_dim=CRITIC_OUTPUT) critic_optim = Adam(critic.parameters(), lr=CRITIC_LR) #Copy the weights to make the targets actor_target = nn.Actor(input_dim=NUM_LOCATIONS, hidden_dim=ACTOR_HIDDEN, output_dim=ACTOR_OUTPUT) critic_target = nn.Critic(state_dim=NUM_LOCATIONS, action_dim=ACTOR_OUTPUT, hidden=CRTIC_HIDDEN, output_dim=CRITIC_OUTPUT) actor_target.load_state_dict(actor.state_dict()) critic_target.load_state_dict(critic.state_dict()) env = environment.Env(locations=NUM_LOCATIONS, actor=actor) env.populate() #Initialize buffer prices_l = [] critic_l = [] for e in range(MAX_EPISODES): for i in range(STEPS): if actor.training != True: actor.train() if critic.training != True: critic.train() actor_target.eval() critic_target.eval() state = env.state() pricing = actor(state) pricing_noised = pricing + EXPLORATION * torch.randn( size=pricing.size()) next_state, reward = env.step(pricing_noised) env.add_memory(state=state, action=pricing_noised, reward=reward, next_state=next_state) states, actions, rewards, nexts = env.observe(batch_size=100) pricing_target = actor_target( states) #+ 10*torch.rand(size=pricing.size()) critic_optim.zero_grad() rewards_tenosr = torch.Tensor(rewards).unsqueeze(dim=1) pred_future_reward = DISCOUNT * critic_target( nexts, pricing_target) y_i = rewards_tenosr + pred_future_reward y = critic(states, actions) critic_loss = F.l1_loss(y_i, y) critic_l.append(reward) critic_loss.backward() critic_optim.step() actor_optim.zero_grad() pricing = actor(states) policy_loss = -critic(states, pricing) - env.wtp_loss(pricing) policy_loss = policy_loss.mean() policy_loss.backward() actor_optim.step() critic.eval() soft_update(actor_target, actor, TAU) soft_update(critic_target, critic, TAU) env.update(next_state) plt.plot(critic_l) plt.show()
def select_action(model, state): # start with e-greedy epsilon = 0.1 if np.random.uniform() > epsilon: return model(state).data.max(1)[1].numpy()[0] else: return np.random.randint(4) np.random.seed(10) envs = [gym.make('small-maze-{}-v0'.format(i)) for i in range(10)] mems = [ReplayMemory() for i in range(10)] dqn = DQN().train() optimizer = Adam(dqn.parameters()) batch_size = 10 gamma = 0.99 episode_steps = [] num_episodes = 100 for i_episode in range(num_episodes): state = env.reset() t = 0 done = False total_reward = 0 while not done and t < 2000: # acting # state_variable = Variable(torch.from_numpy(state[np.newaxis,:,:].astype('float32')))
default=False, metavar='UT', help='load pretrained model (default: False)') args = parser.parse_args() batch_loader = BatchLoader('') parameters = Parameters(batch_loader.max_word_len, batch_loader.max_seq_len, batch_loader.vocab_size) rgan = RGAN(parameters) if args.use_trained: rgan.load_state_dict(t.load('trained_RGAN')) if args.use_cuda: rgan = rgan.cuda() g_optimizer = Adam(rgan.generator.parameters(), args.learning_rate) d_optimizer = Adam(rgan.discriminator.parameters(), args.learning_rate) for iteration in range(args.num_iterations): for _ in range(5): '''Dicriminator forward-loss-backward-update''' z, true_data = batch_loader.input_data(args.batch_size, args.use_cuda, parameters) discriminator_loss, _ = rgan(z, true_data) d_optimizer.zero_grad() discriminator_loss.backward() d_optimizer.step() for p in rgan.discriminator.parameters(): p.data.clamp_(-0.025, 0.025)
def __init__(self, device, cfg): self.device = device self.cfg = cfg self.bs = self.cfg.getint('training', 'batchsize') #Data pipeline self.data = Data(cfg, save=False) ds_train = DS(self.data, cfg) if len(ds_train) != 0: self.loader_train = DataLoader( ds_train, batch_size=self.bs, shuffle=True, drop_last=False, pin_memory=True, num_workers=0, ) else: self.loader_train = [] self.steps_per_epoch = int(len(self.loader_train) / (self.cfg.getint('training', 'n_critic') + 1)) print(len(self.loader_train), self.steps_per_epoch) self.ff = self.data.ff_inp if len(ds_train) != 0: ds_val = DS(self.data, cfg, train=False) if len(ds_val) != 0: self.loader_val = DataLoader( ds_val, batch_size=self.bs, shuffle=True, drop_last=False, pin_memory=True, num_workers=0, ) else: self.loader_val = [] #model self.name = cfg.get('model', 'name') self.cond = cfg.getboolean('model', 'cond') if self.cond and not self.data.pairs: raise Exception('conditional GAN can only be used with pairs of snapshots for both resolutions.') self.out_env = cfg.getboolean('model', 'out_env') self.n_env_mols = cfg.getint('universe', 'n_env_mols') self.feature_dim = self.ff.n_channels """ if self.n_env_mols != 0: self.feature_dim += self.ff_inp.n_atom_chns if self.out_env: self.feature_dim += self.ff_out.n_atom_chns """ self.target_dim = 1 self.critic_dim = self.feature_dim + self.target_dim #if self.cond: #self.critic_dim += self.feature_dim self.z_dim = int(cfg.getint('model', 'noise_dim')) #print(self.ff_inp.n_atom_chns, self.n_input, self.n_out, self.ff_out.n_atoms) self.step = 0 self.epoch = 0 # Make Dirs for saving self.out = OutputHandler( self.name, self.cfg.getint('training', 'n_checkpoints'), self.cfg.get('model', 'output_dir'), ) self.energy = Energy_torch(self.ff, self.device) #self.energy_out = Energy_torch(self.ff_out, self.device) self.n_bins = 64 self.gauss_hist_bond = GaussianHistogram_Dis(bins=self.n_bins, min=0.0, max=0.8, sigma=0.005, ff=self.ff, device=device) self.gauss_hist_angle = GaussianHistogram_Angle(bins=self.n_bins, min=0, max=180, sigma=2.0, ff=self.ff, device=device) self.gauss_hist_dih = GaussianHistogram_Dih(bins=self.n_bins, min=0, max=180, sigma=4.0, ff=self.ff, device=device) self.gauss_hist_nb = GaussianHistogram_Dis(bins=self.n_bins, min=0.0, max=2.0, sigma=0.005, ff=self.ff, device=device) self.ol_weight = cfg.getfloat('prior', 'ol') prior_weights = self.cfg.get('prior', 'weights') self.prior_weights = [float(v) for v in prior_weights.split(",")] prior_schedule = self.cfg.get('prior', 'schedule') try: self.prior_schedule = np.array([0] + [int(v) for v in prior_schedule.split(",")]) except: self.prior_schedule = [0] self.ratio_bonded_nonbonded = cfg.getfloat('prior', 'ratio_bonded_nonbonded') self.prior_mode = cfg.get('prior', 'mode') print(self.prior_mode) #Model selection #if cfg.get('model', 'model_type') == "tiny": # print("Using tiny model") if self.z_dim != 0: self.generator = model.G_tiny_with_noise(z_dim=self.z_dim, n_input=self.feature_dim, n_output=self.target_dim, start_channels=self.cfg.getint('model', 'n_chns'), fac=1, sn=self.cfg.getint('model', 'sn_gen'), device=device) print("Using tiny generator with noise") else: self.generator = model.G_tiny(n_input=self.feature_dim, n_output=self.target_dim, start_channels=self.cfg.getint('model', 'n_chns'), fac=1, sn=self.cfg.getint('model', 'sn_gen'), device=device) print("Using tiny generator without noise") if cfg.getint('grid', 'resolution') == 8: self.critic = model.C_tiny_mbd(in_channels=self.critic_dim, start_channels=self.cfg.getint('model', 'n_chns'), fac=1, sn=self.cfg.getint('model', 'sn_crit'), device=device) print("Using tiny critic with resolution 8") else: self.critic = model.C_tiny16(in_channels=self.critic_dim, start_channels=self.cfg.getint('model', 'n_chns'), fac=1, sn=self.cfg.getint('model', 'sn_crit'), device=device) print("Using tiny critic with resolution 16") self.use_gp = cfg.getboolean('model', 'gp') self.use_ol = cfg.getboolean('model', 'ol') self.use_energy = cfg.getboolean('model', 'energy') self.critic.to(device=device) self.generator.to(device=device) #self.mse.to(device=device) lr_gen = cfg.getfloat('training', 'lr_gen') lr_crit = cfg.getfloat('training', 'lr_crit') #self.opt_generator_pretrain = Adam(self.generator.parameters(), lr=lr_gen, betas=(0, 0.9)) self.opt_generator = Adam(self.generator.parameters(), lr=lr_gen, betas=(0, 0.9)) self.opt_critic = Adam(self.critic.parameters(), lr=lr_crit, betas=(0, 0.9)) self.restored_model = False self.restore_latest_checkpoint()
def train(n_generators, n_discriminators, noise_dim, ngf, ndf, grad_penalty, sampling, mirror_lr, restart, output_dir, data_type, data_source, depth, print_every, eval_device, eval_fid, fused_noise, device, batch_size, n_iter, loss_type, eval_every, D_lr, G_lr, _seed, _run): torch.random.manual_seed(_seed) np.random.seed(_seed) random.seed(_seed) output_dir = output_dir.replace('$WORK', os.environ['WORK']) if not torch.cuda.is_available(): device = 'cpu' if output_dir is None: output_dir = join(_run.observers[0].dir, 'artifacts') else: output_dir = output_dir if not os.path.exists(output_dir): os.makedirs(output_dir) writer = SummaryWriter(log_dir=output_dir) if data_type == 'synthetic': if data_source == '8gmm': dataset, sampler = make_8gmm() elif data_source == '25gmm': dataset, sampler = make_25gmm() else: raise ValueError() dataset.tensors = dataset.tensors.to(device) stat_path = None elif data_type == 'image': dataset, stat_path = make_image_data(data_source) sampler = None else: raise ValueError() data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True) data_loader = infinite_iter(data_loader) def make_generator(): if data_type == 'synthetic': return GeneratorSynthetic(depth=depth, n_filters=ngf) else: if data_source == 'mnist': return GeneratorDCGAN28(in_features=noise_dim, out_channels=1, n_filters=ngf) elif data_source == 'multimnist': return GeneratorDCGAN28(in_features=noise_dim, out_channels=3, n_filters=ngf) elif data_source == 'cifar10': return GeneratorResNet32(n_in=noise_dim, n_out=3, num_filters=ngf) def make_discriminator(): if data_type == 'synthetic': return DiscriminatorSynthetic(depth=depth, n_filters=ndf) else: batch_norm = loss_type != 'wgan-gp' if data_source == 'mnist': return DiscriminatorDCGAN28(in_channels=1, n_filters=ndf, n_targets=1, batch_norm=batch_norm) elif data_source == 'multimnist': return DiscriminatorDCGAN28(in_channels=3, n_filters=ndf, n_targets=1, batch_norm=batch_norm) elif data_source == 'cifar10': return DiscriminatorResNet32(n_in=3, num_filters=ndf) players = {'G': {}, 'D': {}} for G in range(n_generators): players['G'][G] = make_generator().to(device) for D in range(n_discriminators): players['D'][D] = make_discriminator().to(device) for group, these_players in players.items(): for player in these_players.values(): player.train() averagers = { group: {i: ParamAverager(player) for i, player in these_players.items()} for group, these_players in players.items() } optimizers = { group: { i: ExtraOptimizer( Adam(player.parameters(), lr=D_lr if group == 'D' else G_lr, betas=(0.5, 0.9))) for i, player in these_players.items() } for group, these_players in players.items() } log_weights = { 'G': torch.full((n_generators, ), fill_value=1 / n_generators, device=device), 'D': torch.full((n_discriminators, ), fill_value=1 / n_discriminators, device=device) } losses = { 'G': torch.zeros((n_generators, n_discriminators), device=device), 'D': torch.zeros((n_discriminators, n_generators), device=device) } counts = { 'G': torch.zeros((n_generators, n_discriminators), device=device), 'D': torch.zeros((n_discriminators, n_generators), device=device) } last_losses = {'G': None, 'D': None} fixed_data = next(data_loader)[0].to('cpu') if data_type == 'synthetic': fixed_noise = torch.randn(512, noise_dim).to(device) noise_loader = None true_loglike = sampler.log_prob(fixed_data).mean().item() else: fixed_noise = torch.randn(10000, noise_dim).to(device) noise_loader = DataLoader(TensorDataset(fixed_noise), batch_size=batch_size) true_loglike = None scheduler = Scheduler(n_generators, n_discriminators, sampling=sampling) n_gen_upd, n_steps, n_data, n_noise = 0, 0, 0, 0 next_print_step = 0 next_eval_step = 0 if 'checkpoint' in os.listdir(output_dir) and restart: for name in ['checkpoint', 'last_checkpoint']: try: checkpoint = torch.load(join(output_dir, name), map_location=torch.device('cpu')) bundle = joblib.load(join(output_dir, f'{name}_bundle')) except EOFError: print(f'Cannot open {name}') continue else: print(f'Restarting from {join(output_dir, name)}') n_gen_upd = checkpoint['n_gen_upd'] n_steps = checkpoint['n_steps'] n_data = checkpoint['n_data'] n_noise = checkpoint['n_noise'] next_print_step = checkpoint['next_print_step'] next_eval_step = checkpoint['next_eval_step'] for group, these_players in players.items(): for P, player in these_players.items(): players[group][P].load_state_dict( checkpoint['players'][group][P]) optimizers[group][P].load_state_dict( checkpoint['optimizers'][group][P]) averagers[group][P].load_state_dict( checkpoint['averagers'][group][P]) # log_weights = checkpoint['log_weights'][group] losses[group] = checkpoint['past_losses'][group] counts[group] = checkpoint['count_losses'][group] scheduler = bundle['scheduler'] break fake_images = {} for G, generator in players['G'].items(): generator.eval() with torch.no_grad(): fake_images[G] = ((generator(fixed_noise[:8]) + 1) / 2).numpy() joblib.dump(fake_images, 'fake_images.pkl')
def main(): from torch.autograd import Variable from torch.optim import Adam from torchvision import datasets, transforms from torch.autograd import Variable from torch.optim import Adam from torchvision import datasets, transforms reconstruct = False num_classes = 10 lr = [1e-4, 1e-4] batch_size = 16 test_batch_size = 4 # disp_after = 10 num_epochs = 1000 disp_frequency = 1 hidden_size = 128 input_size = 1 num_layers = 1 network = Network(num_classes, hidden_size=hidden_size, input_size=input_size, num_layers=num_layers) model = network.model print model # model.load_state_dict(torch.load('epochs/epoch_327.pt')) model.cuda() print("# parameters:", sum(param.numel() for param in model.parameters())) optimizer = Adam(network.get_lr_list(lr)) # model.parameters(),lr = lr) transformer = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, ))]) train_data = datasets.MNIST('../../data/mnist_data', train=True, transform=transformer, download=True) test_data = datasets.MNIST('../../data/mnist_data', train=False, download=True, transform=transformer) train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(test_data, batch_size=test_batch_size, shuffle=False) model.train() loss_fun = nn.CrossEntropyLoss().cuda() for batch_idx, batch in enumerate(train_loader): data = batch[0] if input_size == 28: data = data.squeeze() elif input_size == 1: data = data.view(data.size(0), data.size(1), -1).transpose(1, 2).contiguous() # print data.size() # raw_input() labels_a = batch[1] break for epoch_num in range(num_epochs): # labels = labels.cuda() labels = Variable(labels_a.cuda()) classes = model(Variable(data).cuda()) loss = loss_fun(classes, labels) optimizer.zero_grad() loss.backward() optimizer.step() torch.nn.utils.clip_grad_norm_(model.parameters(), 1) # print gt_np,pred_np if epoch_num % disp_frequency == 0: _, preds = torch.max(classes, 1) gt_np = labels_a.numpy() pred_np = preds.data.cpu().numpy() accuracy = np.sum(gt_np == pred_np) / float(pred_np.size) print gt_np, pred_np # , torch.norm(model.lstm.grad.data).cpu().numpy()[0],torch.norm(model.classifier.grad.data).cpu().numpy()[0] for p in model.parameters(): print torch.norm(p.grad.data).cpu().numpy() print 'Epoch %d, Loss %.2f, Accuracy %.2f' % (epoch_num, loss, accuracy) if accuracy == 1: raw_input()
def _dqn(env, writers=None): action_repeat = 4 last_timestep = last_frame / action_repeat last_update = (last_timestep - replay_start_size) / update_frequency final_exploration_step = final_exploration_frame / action_repeat n_agents = len(env.agents) n_actions = env.action_spaces['first_0'].n model = model_constructor(env).to(device) optimizer = Adam( model.parameters(), lr=lr, eps=eps ) q = Approximation( model, optimizer, scheduler=CosineAnnealingLR(optimizer, last_update), target=FixedTarget(target_update_frequency), writer=writers['first_0'] ) replay_buffer = ExperienceReplayBuffer( replay_buffer_size, store_device=device, device=device ) def agent_constructor(writer): policy = GreedyPolicy( q, n_actions, epsilon=LinearScheduler( initial_exploration, final_exploration, replay_start_size, final_exploration_step - replay_start_size, name="epsilon", writer=writer ) ) return DeepmindAtariBody( DQN( q, policy, replay_buffer, discount_factor=discount_factor, loss=smooth_l1_loss, minibatch_size=minibatch_size, replay_start_size=replay_start_size, update_frequency=update_frequency, ), lazy_frames=True ) return MultiagentEncoder(IndependentMultiagent({ agent : agent_constructor(writers[agent]) for agent in env.agents }), env.agents, device)
from utils.loss_utils import MarginLoss capsule_loss = MarginLoss(options) elif options.loss_type == 'spread': from utils.loss_utils import SpreadLoss capsule_loss = SpreadLoss(options) elif options.loss_type == 'cross-entropy': capsule_loss = nn.CrossEntropyLoss() if options.add_decoder: from utils.loss_utils import ReconstructionLoss reconst_loss = ReconstructionLoss() optimizer = Adam(capsule_net.parameters(), lr=options.lr, betas=(options.beta1, 0.999)) # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.9) # Attention Regularization l2_loss = nn.MSELoss() ################################## # Load dataset ################################## if options.data_name == 'mnist': from dataset.mnist import MNIST as data os.system('cp {}/dataset/mnist.py {}'.format(BASE_DIR, save_dir)) elif options.data_name == 'fashion_mnist': from dataset.fashion_mnist import FashionMNIST as data os.system('cp {}/dataset/fashion_mnist.py {}'.format(
batch_size = experiment_config[subset]['batch_size'] num_epochs = experiment_config[subset]['num_epochs'] lr = experiment_config[subset]['lr'] train_loader = DataLoader(dataset=train_set, batch_size=batch_size, shuffle=True, collate_fn=collate2) energy_model = MPNN_energy(node_input_dim=74, edge_input_dim=1, edge_hidden_dim=2, node_hidden_dim=10, output_dim=8).to(device) optimizer = Adam(energy_model.parameters(), lr=lr) descriptor = tqdm.trange(num_epochs) train_labels = torch.stack( [train_set.dataset.labels[i] for i in train_set.indices]).to(device) train_mean = torch.mean(train_labels) train_std = torch.std(train_labels) train_meter = Meter(train_mean, train_std) for epoch in descriptor: epoch_loss = 0 for i_batch, batch_data in enumerate(train_loader): indices, protein_graph_bg, ligand_graph_bg, complex_graph_bg, labels = batch_data protein_graph_bg, ligand_graph_bg, complex_graph_bg, labels = protein_graph_bg.to( device), ligand_graph_bg.to(device), complex_graph_bg.to( device), labels.to(device) pred_binding_energy = pred_binding_energy_for_batch(
def __init__( self, env, policy_model, value_model, lr=5e-4, ent_coef=0.01, vf_coef=0.5, ## hyper-parawmeter gamma=0.99, lam=0.95, cliprange=0.2, batch_size=64, value_train_round=200, running_step=2048, running_ep=20, value_regular=0.01, buffer_size=50000, ## decay decay=False, decay_rate=0.9, lstm_enable=False, ## path=None): self.gpu = False self.env = env self.gamma = gamma self.lam = lam self.ent_coef = ent_coef self.vf_coef = vf_coef self.cliprange = cliprange self.value_train_step = value_train_round self.sample_rollout = running_step self.sample_ep = running_ep self.batch_size = batch_size self.lstm_enable = lstm_enable self.replay_buffer = ReplayMemory(buffer_size, other_record=["value", "return"]) self.loss_cal = torch.nn.SmoothL1Loss() self.policy = policy_model if value_model == "shared": self.value = policy_model elif value_model == "copy": self.value = deepcopy(policy_model) else: self.value = value_model self.dist = make_pdtype(env.action_space, policy_model) self.policy_model_optim = Adam(self.policy.parameters(), lr=lr) self.value_model_optim = Adam(self.value.parameters(), lr=lr, weight_decay=value_regular) if decay: self.policy_model_decay_optim = torch.optim.lr_scheduler.ExponentialLR( self.policy_model_optim, decay_rate, last_epoch=-1) self.value_model_decay_optim = torch.optim.lr_scheduler.ExponentialLR( self.value_model_optim, decay_rate, last_epoch=-1) super(PPO_Agent, self).__init__(path) #example_input = Variable(torch.rand((100,)+self.env.observation_space.shape)) #self.writer.add_graph(self.policy, input_to_model=example_input) self.backward_step_show_list = ["pg_loss", "entropy", "vf_loss"] self.backward_ep_show_list = ["pg_loss", "entropy", "vf_loss"] self.training_round = 0 self.running_step = 0 self.record_sample = None self.training_step = 0 self.lstm_enable = True
def train_multiple_epochs(train_dataset, test_dataset, train_u_indices, test_u_indices, model, args, lr_decay_factor, lr_decay_step_size, weight_decay, ARR=0, logger=None, continue_from=None, res_dir=None): rmses = [] # train_loader = DataLoader(train_dataset, batch_size, shuffle=True, num_workers=mp.cpu_count()) # test_loader = DataLoader(test_dataset, batch_size, shuffle=False, num_workers=mp.cpu_count()) model.to(device).reset_parameters() optimizer = Adam(model.parameters(), lr=args.lr, weight_decay=weight_decay) start_epoch = 1 if continue_from is not None: model.load_state_dict( torch.load( os.path.join(res_dir, 'model_checkpoint{}.pth'.format(continue_from)))) optimizer.load_state_dict( torch.load( os.path.join( res_dir, 'optimizer_checkpoint{}.pth'.format(continue_from)))) start_epoch = continue_from + 1 if torch.cuda.is_available(): torch.cuda.synchronize() t_start = time.perf_counter() pbar = tqdm(range(start_epoch, args.epochs + start_epoch)) for epoch in pbar: train_loss = train(model, optimizer, train_loader, device, regression=True, ARR=ARR) rmses.append(eval_rmse(model, test_loader, device)) eval_info = { 'epoch': epoch, 'train_loss': train_loss, 'test_rmse': rmses[-1], } pbar.set_description( 'Epoch {}, train loss {:.6f}, test rmse {:.6f}'.format( *eval_info.values())) if epoch % lr_decay_step_size == 0: for param_group in optimizer.param_groups: param_group['lr'] = lr_decay_factor * param_group['lr'] if logger is not None: logger(eval_info, model, optimizer) if torch.cuda.is_available(): torch.cuda.synchronize() t_end = time.perf_counter() duration = t_end - t_start print('Final Test RMSE: {:.6f}, Duration: {:.6f}'.format( rmses[-1], duration)) return rmses[-1]
def main(args): args['device'] = "cuda" if torch.cuda.is_available() else "cpu" set_random_seed() # Interchangeable with other Dataset if args['dataset'] == 'Tox21': from dgl.data.chem import Tox21 dataset = Tox21() trainset, valset, testset = split_dataset(dataset, args['train_val_test_split']) train_loader = DataLoader(trainset, batch_size=args['batch_size'], collate_fn=collate_molgraphs) val_loader = DataLoader(valset, batch_size=args['batch_size'], collate_fn=collate_molgraphs) test_loader = DataLoader(testset, batch_size=args['batch_size'], collate_fn=collate_molgraphs) if args['pre_trained']: args['num_epochs'] = 0 model = model_zoo.chem.load_pretrained(args['exp']) else: # Interchangeable with other models if args['model'] == 'GCN': model = model_zoo.chem.GCNClassifier( in_feats=args['in_feats'], gcn_hidden_feats=args['gcn_hidden_feats'], classifier_hidden_feats=args['classifier_hidden_feats'], n_tasks=dataset.n_tasks) elif args['model'] == 'GAT': model = model_zoo.chem.GATClassifier( in_feats=args['in_feats'], gat_hidden_feats=args['gat_hidden_feats'], num_heads=args['num_heads'], classifier_hidden_feats=args['classifier_hidden_feats'], n_tasks=dataset.n_tasks) loss_criterion = BCEWithLogitsLoss(pos_weight=torch.tensor( dataset.task_pos_weights).to(args['device']), reduction='none') optimizer = Adam(model.parameters(), lr=args['lr']) stopper = EarlyStopping(patience=args['patience']) model.to(args['device']) for epoch in range(args['num_epochs']): # Train run_a_train_epoch(args, epoch, model, train_loader, loss_criterion, optimizer) # Validation and early stop val_roc_auc = run_an_eval_epoch(args, model, val_loader) early_stop = stopper.step(val_roc_auc, model) print( 'epoch {:d}/{:d}, validation roc-auc score {:.4f}, best validation roc-auc score {:.4f}' .format(epoch + 1, args['num_epochs'], val_roc_auc, stopper.best_score)) if early_stop: break if not args['pre_trained']: stopper.load_checkpoint(model) test_roc_auc = run_an_eval_epoch(args, model, test_loader) print('test roc-auc score {:.4f}'.format(test_roc_auc))
def configure_optimizers(self): if not self.hparams.scheduler: layers_with_parameters_lm = [] for module in self.roberta.modules(): for child in module.children(): if (list(child.children()) == [] and list(child.parameters()) != [] and child.__class__.__name__ != "LayerNorm"): layers_with_parameters_lm.append(child) params_lm = [{ "params": layer.parameters(), "lr": self.hparams.learning_rate_lm * self.hparams.lr_decay**index } for index, layer in enumerate( reversed(layers_with_parameters_lm))] layers_with_parameters_ext = [] for module in self.rnn.modules(): for child in module.children(): if (list(child.children()) == [] and list(child.parameters()) != [] and child.__class__.__name__ != "LayerNorm"): layers_with_parameters_ext.append(child) for module in self.classifier.modules(): for child in module.children(): if (list(child.children()) == [] and list(child.parameters()) != [] and child.__class__.__name__ != "LayerNorm"): layers_with_parameters_ext.append(child) params_ext = [{ "params": layer.parameters(), "lr": self.hparams.learning_rate_ext * self.hparams.lr_decay**index } for index, layer in enumerate( reversed(layers_with_parameters_ext))] return Adam(params_lm + params_ext, lr=self.hparams.learning_rate_ext, weight_decay=self.hparams.weight_decay) else: param_optimizer = list(self.rnn.parameters()) param_optimizer += list(self.classifier.parameters()) optimizer_grouped_parameters = [{ "params": [param for param in param_optimizer] }] optimizer = Adam(optimizer_grouped_parameters, lr=self.hparams.learning_rate_ext, weight_decay=self.hparams.weight_decay) scheduler = ReduceLROnPlateau(optimizer=optimizer, factor=0.5, patience=1, verbose=True) scheduler_dict = { "scheduler": scheduler, "name": "LR-Scheduler", "interval": "epoch", "reduce_on_plateau": True, "monitor": "val_checkpoint_on" } return [optimizer], [scheduler_dict]