def create_optimization(self): if self.args.loss_function == 'FocalLoss': self.loss = FocalLoss(gamma=self.args.gamma) else: self.loss = nn.CrossEntropyLoss() if self.args.cuda: self.loss.cuda() if self.args.classify: self.metric_fc = ArcMarginModel(self.args) self.optimizer = RMSprop(self.model.parameters(), self.args.learning_rate, momentum=self.args.momentum, weight_decay=self.args.weight_decay) else: self.optimizer = RMSprop([{ 'params': self.model.parameters() }, { 'params': self.metric_fc.parameters() }], self.args.learning_rate, momentum=self.args.momentum, weight_decay=self.args.weight_decay)
def __init__(self, mac, scheme, logger, args): self.args = args self.n_agents = args.n_agents self.n_actions = args.n_actions self.mac = mac self.logger = logger self.last_target_update_step = 0 self.critic_training_steps = 0 self.log_stats_t = -self.args.learner_log_interval - 1 self.critic = OffPGCritic(scheme, args) self.mixer = QMixer(args) self.target_critic = copy.deepcopy(self.critic) self.target_mixer = copy.deepcopy(self.mixer) self.agent_params = list(mac.parameters()) self.critic_params = list(self.critic.parameters()) self.mixer_params = list(self.mixer.parameters()) self.params = self.agent_params + self.critic_params self.c_params = self.critic_params + self.mixer_params self.agent_optimiser = RMSprop(params=self.agent_params, lr=args.lr) self.critic_optimiser = RMSprop(params=self.critic_params, lr=args.lr) self.mixer_optimiser = RMSprop(params=self.mixer_params, lr=args.lr) print('Mixer Size: ') print(get_parameters_num(list(self.c_params)))
def create_optimizer(args: ModelConfigBase, model: torch.nn.Module) -> Optimizer: """ Creates a torch optimizer for the given model. :param args: The arguments object with attributes used to create the optimizer_type. :param model: The DataParallel object representing the network. :return: An instance of torch.optim.Optimizer """ # Select optimizer type if args.optimizer_type in [OptimizerType.Adam, OptimizerType.AMSGrad]: return torch.optim.Adam( model.parameters(), args.l_rate, args.adam_betas, args.opt_eps, args.weight_decay, amsgrad=args.optimizer_type == OptimizerType.AMSGrad) elif args.optimizer_type == OptimizerType.SGD: return torch.optim.SGD(model.parameters(), args.l_rate, args.momentum, weight_decay=args.weight_decay) elif args.optimizer_type == OptimizerType.RMSprop: return RMSprop(model.parameters(), args.l_rate, args.rms_alpha, args.opt_eps, args.weight_decay, args.momentum) else: raise NotImplementedError( f"Optimizer type {args.optimizer_type.value} is not implemented")
def __init__(self, config, n_Feature, n_Action): self.device = config["device"] self.lr = config["lr"] self.lr_step_size = config["lr_decay_step"] self.lr_gamma = config["lr_decay_gamma"] self.lr_last_epoch = config["lr_last_epoch"] self.target_net_update_freq = config["target_net_update_freq"] self.experience_replay_size = config["experience_replay_size"] self.batch_size = config["batch_size"] self.discount = config["discount"] self.num_feature = n_Feature self.num_action = n_Action # initialize the predict net and target net in the agent self.predict_net = DNN(input_shape=n_Feature, num_actions=n_Action) self.target_net = DNN(input_shape=n_Feature, num_actions=n_Action) self.predict_net = self.predict_net.to(self.device) self.target_net = self.target_net.to(self.device) self.predict_net.apply(weigth_init) self.target_net.apply(weigth_init) self.target_net.load_state_dict(self.predict_net.state_dict()) self.optimizer = RMSprop(self.predict_net.parameters(), lr=self.lr, momentum=0.95, eps=0.01) self.lr_schduler = StepLR(optimizer=self.optimizer, step_size=self.lr_step_size, gamma=self.lr_gamma) # self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr) # initialize the experience replay buffer self.memory = ReplayMemory(entry_size=n_Feature, memory_size=self.experience_replay_size, batch_size=self.batch_size)
def configure_optimizers(self): eps = 1e-2 / float(self.hparams.batch_size)**2 if self.hparams.optimizer_type == "RMSprop": optimizer = RMSprop(self.parameters(), lr=self.hparams.learning_rate, momentum=0.9, eps=eps, weight_decay=self.hparams.weight_decay) elif self.hparams.optimizer_type == "RAdam": optimizer = RAdam(self.parameters(), lr=self.hparams.learning_rate, eps=eps, weight_decay=self.hparams.weight_decay) elif self.hparams.optimizer_type == "Adam": optimizer = Adam(self.parameters(), lr=self.hparams.learning_rate, eps=eps, weight_decay=self.hparams.weight_decay) else: raise ValueError("Unknown optimizer type.") if self.hparams.look_ahead: optimizer = LookAhead(optimizer, k=self.hparams.look_ahead_k, alpha=self.hparams.look_ahead_alpha) if not self.hparams.use_lr_scheduler: return optimizer scheduler = ExponentialLR(optimizer=optimizer, gamma=self.hparams.lr_scheduler_decay_rate) return [optimizer], [scheduler]
def configure_optimizers(self): lr = self.cfg.optimizer.learning_rate weight_decay = self.cfg.optimizer.weight_decay eps = 1e-2 / float(self.cfg.data_loader.batch_size)**2 if self.cfg.optimizer.type == "RMSprop": optimizer = RMSprop(self.parameters(), lr=lr, momentum=self.cfg.optimizer.momentum, eps=eps, weight_decay=weight_decay) elif self.cfg.optimizer.type == "RAdam": optimizer = RAdam(self.parameters(), lr=lr, eps=eps, weight_decay=weight_decay) elif self.cfg.optimizer.type == "Adam": optimizer = Adam(self.parameters(), lr=lr, eps=eps, weight_decay=weight_decay) else: raise ValueError("Unknown optimizer type.") if self.cfg.meta_optimizer.look_ahead: optimizer = LookAhead(optimizer, k=self.cfg.optimizer.look_ahead_k, alpha=self.cfg.optimizer.look_ahead_alpha) if not self.cfg.lr_scheduler.active: return optimizer scheduler = ExponentialLR(optimizer=optimizer, gamma=self.cfg.lr_scheduler.decay_rate) return [optimizer], [scheduler]
def __init__(self, input_shape, action_space, seed, device, model, gamma, alpha, tau, batch_size,update, replay, buffer_size, env, decay = 200, path = 'model',num_epochs= 0, max_step = 50000, learn_interval = 20): '''Initialise a DQNAgent Object buffer_size : size of replay buffer to sample from gamma : discount rate alpha : learn rate replay. : after which replay buffer loading to be started update : update interval of model parameters every x instances of back propagation replay. : after which replay buffer loading to be started learn_interval: tick for learning rate ''' super(DQNAgent,self).__init__( input_shape ,action_space ,seed ,device,model, gamma, alpha, tau, batch_size, max_step, env,num_epochs ,path) self.buffer_size = buffer_size self.update = update self.replay = replay self.interval = learn_interval # Q-Network self.policy_net = self.model(input_shape, action_space).to(self.device) self.target_net = self.model(input_shape, action_space).to(self.device) self.optimiser = RMSprop(self.policy_net.parameters(), lr=self.alpha) # Replay Memory self.memory = ReplayMemory(self.buffer_size, self.batch_size, self.seed, self.device) # Timestep self.t_step = 0 self.l_step = 0 self.EPSILON_START = 1.0 self.EPSILON_FINAL = 0.02 self.EPS_DECAY = decay self.epsilon_delta = lambda frame_idx: self.EPSILON_FINAL + (self.EPSILON_START - self.EPSILON_FINAL) * exp(-1. * frame_idx / self.EPS_DECAY)
def set_mlp_optim(self): if self.optimizer == "Adam": return torch.optim.Adam(self.net.parameters(), lr=self.lr) elif self.optimizer == "SGD": return torch.optim.SGD(self.net.parameters(), lr=self.lr) elif self.optimizer == "rmsprop": return RMSprop(self.net.parameters(), lr=self.lr)
def configure_optimizers(self): optim = RMSprop(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay) scheduler = torch.optim.lr_scheduler.StepLR(optim, self.hparams.lr_step_size, self.hparams.lr_gamma) return [optim], [scheduler]
def train_g(self, epoch_num, mode='Adam', dataname='MNIST', logname='MNIST'): print(mode) if mode == 'SGD': g_optimizer = optim.SGD(self.G.parameters(), lr=self.lr, weight_decay=self.weight_decay) self.writer_init(logname=logname, comments='SGD-%.3f_%.5f' % (self.lr, self.weight_decay)) elif mode == 'Adam': g_optimizer = optim.Adam(self.G.parameters(), lr=self.lr, weight_decay=self.weight_decay, betas=(0.5, 0.999)) self.writer_init(logname=logname, comments='ADAM-%.3f_%.5f' % (self.lr, self.weight_decay)) elif mode == 'RMSProp': g_optimizer = RMSprop(self.G.parameters(), lr=self.lr, weight_decay=self.weight_decay) self.writer_init(logname=logname, comments='RMSProp-%.3f_%.5f' % (self.lr, self.weight_decay)) timer = time.time() for e in range(epoch_num): z = torch.randn((self.batchsize, self.z_dim), device=self.device) ## changed fake_x = self.G(z) d_fake = self.D(fake_x) # G_loss = g_loss(d_fake) G_loss = self.criterion( d_fake, torch.ones(d_fake.shape, device=self.device)) g_optimizer.zero_grad() zero_grad(self.D.parameters()) G_loss.backward() g_optimizer.step() gd = torch.norm(torch.cat( [p.grad.contiguous().view(-1) for p in self.D.parameters()]), p=2) gg = torch.norm(torch.cat( [p.grad.contiguous().view(-1) for p in self.G.parameters()]), p=2) self.plot_param(G_loss=G_loss) self.plot_grad(gd=gd, gg=gg) if self.count % self.show_iter == 0: self.show_info(timer=time.time() - timer, D_loss=G_loss) timer = time.time() self.count += 1 if self.count % 5000 == 0: self.save_checkpoint('fixD_%s-%.5f_%d.pth' % (mode, self.lr, self.count), dataset=dataname) self.writer.close()
def get_optimizer_for_model(optimizer_dict: dict, model: Module) -> Optimizer: optimizer_type = optimizer_dict.pop("type") if optimizer_type == "adam": optimizer = Adam(model.parameters(), **optimizer_dict) elif optimizer_type == "rmsprop": optimizer = RMSprop(model.parameters(), **optimizer_dict) else: raise NotImplementedError( 'Unknown optimizer type "{}"'.format(optimizer_type)) return optimizer
def create_optimization(self): self.loss = nn.CrossEntropyLoss() if self.args.cuda: self.loss.cuda() self.optimizer = RMSprop(self.model.parameters(), self.args.learning_rate, momentum=self.args.momentum, weight_decay=self.args.weight_decay)
def create_optimizer(config: DeepLearningConfig, parameters: Iterator[Parameter]) -> torch.optim.Optimizer: # Select optimizer type if config.optimizer_type in [OptimizerType.Adam, OptimizerType.AMSGrad]: return torch.optim.Adam(parameters, config.l_rate, config.adam_betas, config.opt_eps, config.weight_decay, amsgrad=config.optimizer_type == OptimizerType.AMSGrad) elif config.optimizer_type == OptimizerType.SGD: return torch.optim.SGD(parameters, config.l_rate, config.momentum, weight_decay=config.weight_decay) elif config.optimizer_type == OptimizerType.RMSprop: return RMSprop(parameters, config.l_rate, config.rms_alpha, config.opt_eps, config.weight_decay, config.momentum) else: raise NotImplementedError(f"Optimizer type {config.optimizer_type.value} is not implemented")
def get_optimizer(model: nn.Module, optim: str, lr: float) -> Optimizer: """ Return the optimizer that corresponds to string optim. Add the parameters from model and set learning rate to lr :param model: model to get the parameters from :param optim: name of the optimizer :param lr: learning rate to use in the optimizer :return: """ if optim == "adagrad": return Adagrad(model.parameters(), lr=lr) elif optim == "sgd": return SGD(model.parameters(), lr=lr) elif optim == "rmsprop": return RMSprop(model.parameters(), lr=lr) elif optim == "adam": return Adam(model.parameters(), lr=lr) else: raise ValueError("Invalid optimizer")
def create_optimizer(self) -> None: """ Creates a torch optimizer for the given model, and stores it as an instance variable in the current object. """ # Make sure model is created before we create optimizer if self._model is None: raise ValueError("Model checkpoint must be created before optimizer checkpoint can be loaded.") # Select optimizer type if self.config.optimizer_type in [OptimizerType.Adam, OptimizerType.AMSGrad]: self._optimizer = torch.optim.Adam(self._model.parameters(), self.config.l_rate, self.config.adam_betas, self.config.opt_eps, self.config.weight_decay, amsgrad=self.config.optimizer_type == OptimizerType.AMSGrad) elif self.config.optimizer_type == OptimizerType.SGD: self._optimizer = torch.optim.SGD(self._model.parameters(), self.config.l_rate, self.config.momentum, weight_decay=self.config.weight_decay) elif self.config.optimizer_type == OptimizerType.RMSprop: self._optimizer = RMSprop(self._model.parameters(), self.config.l_rate, self.config.rms_alpha, self.config.opt_eps, self.config.weight_decay, self.config.momentum) else: raise NotImplementedError(f"Optimizer type {self.config.optimizer_type.value} is not implemented")
def get_optimizer(optimizer_name: str, parameters, learning_rate: float, weight_decay=1e-5, **kwargs): if optimizer_name.lower() == 'sgd': return SGD(parameters, learning_rate, momentum=0.9, nesterov=True, weight_decay=weight_decay, **kwargs) if optimizer_name.lower() == 'adam': return Adam(parameters, learning_rate, weight_decay=weight_decay, eps=1e-3, # As Jeremy suggests **kwargs) if optimizer_name.lower() == 'rms': return RMSprop(parameters, learning_rate, weight_decay=weight_decay, **kwargs) if optimizer_name.lower() == 'adamw': return AdamW(parameters, learning_rate, weight_decay=weight_decay, eps=1e-3, **kwargs) if optimizer_name.lower() == 'radam': return RAdam(parameters, learning_rate, weight_decay=weight_decay, eps=1e-3, # As Jeremy suggests **kwargs) if optimizer_name.lower() == 'ranger': return Ranger(parameters, learning_rate, weight_decay=weight_decay, **kwargs) if optimizer_name.lower() == 'qhadamw': return QHAdamW(parameters, learning_rate, weight_decay=weight_decay, **kwargs) if optimizer_name.lower() == 'lamb': return Lamb(parameters, learning_rate, weight_decay=weight_decay, **kwargs) if optimizer_name.lower() == 'lamb': return Lamb(parameters, learning_rate, weight_decay=weight_decay, **kwargs) raise ValueError("Unsupported optimizer name " + optimizer_name)
def create_optim(params, args=None, optim_name=None, lr=None): if args is None: assert optim_name is not None and lr is not None else: assert optim_name is None and lr is None optim_name = args.optim lr = args.lr if optim_name == 'sgd': return SGD(params, lr=lr, momentum=0.9, weight_decay=5e-4, nesterov=True) elif optim_name == 'adam': return Adam(params, lr=lr, weight_decay=0) elif optim_name == 'adamw': return AdamW(params, lr=lr, weight_decay=1e-2) elif optim_name == 'amsgrad': return Adam(params, lr=lr, weight_decay=0, amsgrad=True) elif optim_name == 'rmsprop': return RMSprop(params, lr=lr, momentum=0.9, weight_decay=0) else: raise NotImplementedError( 'Unsupported optimizer_memory: {}'.format(optim_name))
def train_gd(self, epoch_num, mode='Adam', dataname='MNIST', logname='MNIST', loss_type='JSD'): print(mode) if mode == 'SGD': d_optimizer = optim.SGD(self.D.parameters(), lr=self.lr, weight_decay=self.weight_decay) g_optimizer = optim.SGD(self.G.parameters(), lr=self.lr, weight_decay=self.weight_decay) elif mode == 'Adam': d_optimizer = optim.Adam(self.D.parameters(), lr=self.lr, weight_decay=self.weight_decay, betas=(0.5, 0.999)) g_optimizer = optim.Adam(self.G.parameters(), lr=self.lr, weight_decay=self.weight_decay, betas=(0.5, 0.999)) elif mode == 'RMSProp': d_optimizer = RMSprop(self.D.parameters(), lr=self.lr, weight_decay=self.weight_decay) g_optimizer = RMSprop(self.G.parameters(), lr=self.lr, weight_decay=self.weight_decay) self.writer_init(logname=logname, comments='%s-%.3f_%.5f' % (mode, self.lr, self.weight_decay)) self.iswriter.writeheader() timer = time.time() start = time.time() for e in range(epoch_num): for real_x in self.dataloader: real_x = real_x[0].to(self.device) d_real = self.D(real_x) z = torch.randn((self.batchsize, self.z_dim), device=self.device) ## changed (shape) fake_x = self.G(z) d_fake = self.D(fake_x.detach()) if loss_type == 'JSD': loss = self.criterion(d_real, torch.ones(d_real.shape, device=self.device)) + \ self.criterion(d_fake, torch.zeros(d_fake.shape, device=self.device)) else: loss = d_fake.mean() - d_real.mean() # D_loss = gan_loss(d_real, d_fake) # D_loss = self.criterion(d_real, torch.ones(d_real.shape, device=self.device)) + \ # self.criterion(d_fake, torch.zeros(d_fake.shape, device=self.device)) D_loss = loss + self.l2penalty() d_optimizer.zero_grad() D_loss.backward() d_optimizer.step() z = torch.randn((self.batchsize, self.z_dim), device=self.device) ## changed fake_x = self.G(z) d_fake = self.D(fake_x) # G_loss = g_loss(d_fake) if loss_type == 'JSD': G_loss = self.criterion( d_fake, torch.ones(d_fake.shape, device=self.device)) else: G_loss = -d_fake.mean() g_optimizer.zero_grad() G_loss.backward() g_optimizer.step() gd = torch.norm(torch.cat([ p.grad.contiguous().view(-1) for p in self.D.parameters() ]), p=2) gg = torch.norm(torch.cat([ p.grad.contiguous().view(-1) for p in self.G.parameters() ]), p=2) self.plot_param(D_loss=D_loss, G_loss=G_loss) self.plot_grad(gd=gd, gg=gg) self.plot_d(d_real, d_fake) if self.count % self.show_iter == 0: self.show_info(timer=time.time() - timer, D_loss=D_loss, G_loss=G_loss) timer = time.time() self.count += 1 if self.count % 2000 == 0: is_mean, is_std = self.get_inception_score(batch_num=500) print(is_mean, is_std) self.iswriter.writerow({ 'iter': self.count, 'is_mean': is_mean, 'is_std': is_std, 'time': time.time() - start }) self.save_checkpoint('%s-%.5f_%d.pth' % (mode, self.lr, self.count), dataset=dataname) self.writer.close() self.save_checkpoint('DIM64%s-%.5f_%d.pth' % (mode, self.lr, self.count), dataset=dataname)
def main(args): # Select the hardware device to use for inference. if torch.cuda.is_available(): device = torch.device('cuda', torch.cuda.current_device()) torch.backends.cudnn.benchmark = True else: device = torch.device('cpu') # Disable gradient calculations by default. torch.set_grad_enabled(False) # create checkpoint dir os.makedirs(args.checkpoint, exist_ok=True) if args.arch == 'hg1': model = hg1(pretrained=False) elif args.arch == 'hg2': model = hg2(pretrained=False) elif args.arch == 'hg8': model = hg8(pretrained=False) else: raise Exception('unrecognised model architecture: ' + args.arch) model = DataParallel(model).to(device) optimizer = RMSprop(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) best_acc = 0 # optionally resume from a checkpoint if args.resume: assert os.path.isfile(args.resume) print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_acc = checkpoint['best_acc'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) logger = Logger(os.path.join(args.checkpoint, 'log.txt'), resume=True) else: logger = Logger(os.path.join(args.checkpoint, 'log.txt')) logger.set_names( ['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc']) # create data loader train_dataset = Mpii(args.image_path, is_train=True) train_loader = DataLoader(train_dataset, batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=True) val_dataset = Mpii(args.image_path, is_train=False) val_loader = DataLoader(val_dataset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=True) # train and eval lr = args.lr for epoch in trange(args.start_epoch, args.epochs, desc='Overall', ascii=True): lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule, args.gamma) # train for one epoch train_loss, train_acc = do_training_epoch(train_loader, model, device, Mpii.DATA_INFO, optimizer, acc_joints=Mpii.ACC_JOINTS) # evaluate on validation set valid_loss, valid_acc, predictions = do_validation_epoch( val_loader, model, device, Mpii.DATA_INFO, False, acc_joints=Mpii.ACC_JOINTS) # print metrics tqdm.write( f'[{epoch + 1:3d}/{args.epochs:3d}] lr={lr:0.2e} ' f'train_loss={train_loss:0.4f} train_acc={100 * train_acc:0.2f} ' f'valid_loss={valid_loss:0.4f} valid_acc={100 * valid_acc:0.2f}') # append logger file logger.append( [epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc]) logger.plot_to_file(os.path.join(args.checkpoint, 'log.svg'), ['Train Acc', 'Val Acc']) # remember best acc and save checkpoint is_best = valid_acc > best_acc best_acc = max(valid_acc, best_acc) save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, predictions, is_best, checkpoint=args.checkpoint, snapshot=args.snapshot) logger.close()
def main_worker(gpu, ngpus_per_node, args): args.gpu = gpu if args.gpu is not None: print(f"Use GPU: {args.gpu} for training!") if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) # create model if 'g' in args.generator_arch: if args.pretrained: generator = Generator.from_pretrained(args.generator_arch) print(f"=> using pre-trained model '{args.generator_arch}'") else: print(f"=> creating model '{args.generator_arch}'") generator = Generator.from_name(args.generator_arch) else: warnings.warn( 'You have chosen a specific model architecture. This will ' 'default use MNIST model architecture!') if args.pretrained: generator = Generator.from_pretrained('g-mnist') print(f"=> using pre-trained model `g-mnist`") else: print(f"=> creating model `g-mnist`") generator = Generator.from_name('g-mnist') if 'd' in args.discriminator_arch: if args.pretrained: discriminator = Discriminator.from_pretrained( args.discriminator_arch) print(f"=> using pre-trained model '{args.discriminator_arch}'") else: print(f"=> creating model `{args.discriminator_arch}`") discriminator = Discriminator.from_name(args.discriminator_arch) else: warnings.warn( 'You have chosen a specific model architecture. This will ' 'default use MNIST model architecture!') if args.pretrained: discriminator = Discriminator.from_pretrained('d-mnist') print(f"=> using pre-trained model `d-mnist`") else: print(f"=> creating model `d-mnist`") discriminator = Discriminator.from_name('d-mnist') if args.distributed: # For multiprocessing distributed, DistributedDataParallel constructor # should always set the single device scope, otherwise, # DistributedDataParallel will use all available devices. if args.gpu is not None: torch.cuda.set_device(args.gpu) generator.cuda(args.gpu) discriminator.cuda(args.gpu) # When using a single GPU per process and per # DistributedDataParallel, we need to divide the batch size # ourselves based on the total number of GPUs we have args.batch_size = int(args.batch_size / ngpus_per_node) args.workers = int(args.workers / ngpus_per_node) generator = torch.nn.parallel.DistributedDataParallel( generator, device_ids=[args.gpu]) discriminator = torch.nn.parallel.DistributedDataParallel( discriminator, device_ids=[args.gpu]) else: generator.cuda() discriminator.cuda() # DistributedDataParallel will divide and allocate batch_size to all # available GPUs if device_ids are not set generator = torch.nn.parallel.DistributedDataParallel(generator) discriminator = torch.nn.parallel.DistributedDataParallel( discriminator) elif args.gpu is not None: torch.cuda.set_device(args.gpu) generator = generator.cuda(args.gpu) discriminator = discriminator.cuda(args.gpu) else: # DataParallel will divide and allocate batch_size to all available # GPUs generator = torch.nn.DataParallel(generator).cuda() discriminator = torch.nn.DataParallel(discriminator).cuda() optimizerG = RMSprop(generator.parameters(), lr=args.lr) optimizerD = RMSprop(discriminator.parameters(), lr=args.lr) # optionally resume from a checkpoint if args.netG: if os.path.isfile(args.netG): print(f"=> loading checkpoint `{args.netG}`") state_dict = torch.load(args.netG) generator.load_state_dict(state_dict) compress_model(state_dict, filename=args.netG, model_arch=args.generator_arch) print(f"=> loaded checkpoint `{args.netG}`") else: print(f"=> no checkpoint found at `{args.netG}`") if args.netD: if os.path.isfile(args.netD): print(f"=> loading checkpoint `{args.netD}`") state_dict = torch.load(args.netD) discriminator.load_state_dict(state_dict) compress_model(state_dict, filename=args.netD, model_arch=args.discriminator_arch) print(f"=> loaded checkpoint `{args.netD}`") else: print(f"=> no checkpoint found at '{args.netD}'") cudnn.benchmark = True if args.name == 'mnist': dataset = datasets.MNIST(root=args.dataroot, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, ), (0.5, )), ])) elif args.name == 'fmnist': dataset = datasets.FashionMNIST(root=args.dataroot, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, ), (0.5, )), ])) else: warnings.warn('You have chosen a specific dataset. This will ' 'default use MNIST dataset!') dataset = datasets.MNIST(root=args.dataroot, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, ), (0.5, )), ])) dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=int(args.workers)) if args.evaluate: validate(generator, args) return for epoch in range(args.start_epoch, args.epochs): # train for one epoch train(dataloader, generator, discriminator, optimizerG, optimizerD, epoch, args) if not args.multiprocessing_distributed or ( args.multiprocessing_distributed and args.rank % ngpus_per_node == 0): # do checkpointing torch.save(generator.state_dict(), f"{args.outf}/netG_epoch_{epoch}.pth") torch.save(discriminator.state_dict(), f"{args.outf}/netD_epoch_{epoch}.pth")
def model_train(self, epoch_offset=0): create_dir(MODEL_SAVE_PATH) loss_for_regression = MSELoss() img_coors_json = read_json_file(BBOX_XYWH_JSON_PATH) optimizer = RMSprop(self.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM) # optimizer = Adam(self.parameters(), lr=LEARNING_RATE) # optimizer = SGD(self.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM) scheduler = StepLR(optimizer, step_size=SCHEDULER_STEP, gamma=SCHEDULER_GAMMA) for epoch in range(EPOCHS): epoch_loss = 0.0 scheduler.step(epoch) LOGGER.debug('Epoch: %s, Current Learning Rate: %s', str(epoch + epoch_offset), str(scheduler.get_lr())) for image, coors in img_coors_json.items(): path_of_image = NORMALISED_IMAGES_PATH + image path_of_image = path_of_image.replace('%', '_') img = cv2.imread(path_of_image) img = torch.tensor(img).float().permute(2, 0, 1).unsqueeze(0) img = img.to(self.device) predicted_width, predicted_height, predicted_midpoint = self.forward( img) #all are scaled mp_x = coors[0][0] mp_y = coors[0][1] mp = torch.cat((torch.tensor([[mp_x]]).to( self.device), torch.tensor([[mp_y]]).to(self.device)), dim=1).float() w = coors[0][2] h = coors[0][3] loss1 = loss_for_regression( predicted_height, torch.tensor([[h]]).float().to(self.device)) loss2 = loss_for_regression( predicted_width, torch.tensor([[w]]).float().to(self.device)) loss3 = loss_for_regression(predicted_midpoint, mp.to(self.device)) loss = loss1 + loss2 + loss3 / 2 optimizer.zero_grad() loss.backward() clip_grad_norm(self.parameters(), 0.5) optimizer.step() epoch_loss = epoch_loss + loss.item() if epoch % 5 == 0: print('epoch: ' + str(epoch) + ' ' + 'loss: ' + str(epoch_loss)) if epoch % EPOCH_SAVE_INTERVAL == 0: print('saving') torch.save( self.state_dict(), MODEL_SAVE_PATH + 'model_epc_' + str(epoch + epoch_offset) + '.pt') torch.save( self.state_dict(), MODEL_SAVE_PATH + 'model_epc_' + str(epoch + epoch_offset) + '.pt')
def update_progess_bar(pbar, params: dict, f=0.99): for param_name, value in params.items(): if "running" in param_name: value = f * pbar.postfix[0][param_name] + (1 - f) * value pbar.postfix[0][param_name] = round(value, 2) pbar.update() def train_agent( rule_agent: RuleBasedAgent, agent: DialogManagerAgent, dialog_env: DialogEnv, train_steps=3_000, batch_size=32, ): optimizer = RMSprop(agent.parameters(), lr=1e-2) warmup_experience_iterator = iter( experience_generator(rule_agent, dialog_env, max_it=1000)) experience_iterator = iter(experience_generator(agent, dialog_env)) exp_it = itertools.chain( *[warmup_experience_iterator, experience_iterator]) agent.exploration_rate = 1.0 min_eps = 0.01 exploration_decay = np.exp(np.log(min_eps) / train_steps) with tqdm(postfix=[{"running_reward": 0.0}]) as pbar: for it in range(train_steps): with torch.no_grad(): agent.eval()
def train_sim(epoch_num=10, optim_type='ACGD', startPoint=None, start_n=0, z_dim=128, batchsize=64, l2_penalty=0.0, momentum=0.0, log=False, loss_name='WGAN', model_name='dc', model_config=None, data_path='None', show_iter=100, logdir='test', dataname='CIFAR10', device='cpu', gpu_num=1): lr_d = 1e-4 lr_g = 1e-4 dataset = get_data(dataname=dataname, path=data_path) dataloader = DataLoader(dataset=dataset, batch_size=batchsize, shuffle=True, num_workers=4) D, G = get_model(model_name=model_name, z_dim=z_dim, configs=model_config) D.apply(weights_init_d).to(device) G.apply(weights_init_g).to(device) optim_d = RMSprop(D.parameters(), lr=lr_d) optim_g = RMSprop(G.parameters(), lr=lr_g) if startPoint is not None: chk = torch.load(startPoint) D.load_state_dict(chk['D']) G.load_state_dict(chk['G']) optim_d.load_state_dict(chk['d_optim']) optim_g.load_state_dict(chk['g_optim']) print('Start from %s' % startPoint) if gpu_num > 1: D = nn.DataParallel(D, list(range(gpu_num))) G = nn.DataParallel(G, list(range(gpu_num))) timer = time.time() count = 0 if 'DCGAN' in model_name: fixed_noise = torch.randn((64, z_dim, 1, 1), device=device) else: fixed_noise = torch.randn((64, z_dim), device=device) for e in range(epoch_num): print('======Epoch: %d / %d======' % (e, epoch_num)) for real_x in dataloader: real_x = real_x[0].to(device) d_real = D(real_x) if 'DCGAN' in model_name: z = torch.randn((d_real.shape[0], z_dim, 1, 1), device=device) else: z = torch.randn((d_real.shape[0], z_dim), device=device) fake_x = G(z) d_fake = D(fake_x) loss = get_loss(name=loss_name, g_loss=False, d_real=d_real, d_fake=d_fake, l2_weight=l2_penalty, D=D) D.zero_grad() G.zero_grad() loss.backward() optim_d.step() optim_g.step() if count % show_iter == 0: time_cost = time.time() - timer print('Iter :%d , Loss: %.5f, time: %.3fs' % (count, loss.item(), time_cost)) timer = time.time() with torch.no_grad(): fake_img = G(fixed_noise).detach() path = 'figs/%s_%s/' % (dataname, logdir) if not os.path.exists(path): os.makedirs(path) vutils.save_image(fake_img, path + 'iter_%d.png' % (count + start_n), normalize=True) save_checkpoint( path=logdir, name='%s-%s%.3f_%d.pth' % (optim_type, model_name, lr_g, count + start_n), D=D, G=G, optimizer=optim_d, g_optimizer=optim_g) if wandb and log: wandb.log({ 'Real score': d_real.mean().item(), 'Fake score': d_fake.mean().item(), 'Loss': loss.item() }) count += 1
def configure_optimizers(self): optimizer = RMSprop(self.parameters(), lr=self.configs.base_lr, momentum=self.configs.momentum, weight_decay=self.configs.weight_decay) return optimizer
def main(args): """Train/ Cross validate for data source = YogiDB.""" # Create data loader """Generic(data.Dataset)(image_set, annotations, is_train=True, inp_res=256, out_res=64, sigma=1, scale_factor=0, rot_factor=0, label_type='Gaussian', rgb_mean=RGB_MEAN, rgb_stddev=RGB_STDDEV).""" annotations_source = 'basic-thresholder' # Get the data from yogi db_obj = YogiDB(config.db_url) imageset = db_obj.get_filtered(ImageSet, name=args.image_set_name) annotations = db_obj.get_annotations(image_set_name=args.image_set_name, annotation_source=annotations_source) pts = torch.Tensor(annotations[0]['joint_self']) num_classes = pts.size(0) crop_size = 512 if args.crop: crop_size = args.crop crop = True else: crop = False # Using the default RGB mean and std dev as 0 RGB_MEAN = torch.as_tensor([0.0, 0.0, 0.0]) RGB_STDDEV = torch.as_tensor([0.0, 0.0, 0.0]) dataset = Generic(image_set=imageset, inp_res=args.inp_res, out_res=args.out_res, annotations=annotations, mode=args.mode, crop=crop, crop_size=crop_size, rgb_mean=RGB_MEAN, rgb_stddev=RGB_STDDEV) train_dataset = dataset train_dataset.is_train = True train_loader = DataLoader(train_dataset, batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=True) val_dataset = dataset val_dataset.is_train = False val_loader = DataLoader(val_dataset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=True) # Select the hardware device to use for inference. if torch.cuda.is_available(): device = torch.device('cuda', torch.cuda.current_device()) torch.backends.cudnn.benchmark = True else: device = torch.device('cpu') # Disable gradient calculations by default. torch.set_grad_enabled(False) # create checkpoint dir os.makedirs(args.checkpoint, exist_ok=True) if args.arch == 'hg1': model = hg1(pretrained=False, num_classes=num_classes) elif args.arch == 'hg2': model = hg2(pretrained=False, num_classes=num_classes) elif args.arch == 'hg8': model = hg8(pretrained=False, num_classes=num_classes) else: raise Exception('unrecognised model architecture: ' + args.model) model = DataParallel(model).to(device) if args.optimizer == "Adam": optimizer = Adam(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) else: optimizer = RMSprop(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) best_acc = 0 # optionally resume from a checkpoint title = args.data_identifier + ' ' + args.arch if args.resume: assert os.path.isfile(args.resume) print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint['epoch'] best_acc = checkpoint['best_acc'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})" .format(args.resume, checkpoint['epoch'])) logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title) logger.set_names(['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc']) # train and eval lr = args.lr for epoch in range(args.start_epoch, args.epochs): lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule, args.gamma) print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr)) # train for one epoch train_loss, train_acc = do_training_epoch(train_loader, model, device, optimizer) # evaluate on validation set if args.debug == 1: valid_loss, valid_acc, predictions, validation_log = do_validation_epoch(val_loader, model, device, False, True, os.path.join(args.checkpoint, 'debug.csv'), epoch + 1) else: valid_loss, valid_acc, predictions, _ = do_validation_epoch(val_loader, model, device, False) # append logger file logger.append([epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc]) # remember best acc and save checkpoint is_best = valid_acc > best_acc best_acc = max(valid_acc, best_acc) save_checkpoint({ 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, predictions, is_best, checkpoint=args.checkpoint, snapshot=args.snapshot) logger.close() logger.plot(['Train Acc', 'Val Acc']) savefig(os.path.join(args.checkpoint, 'log.eps'))
def create_optimizer(optimizer_config, model, master_params=None): if optimizer_config.get("classifier_lr", -1) != -1: # Separate classifier parameters from all others net_params = [] classifier_params = [] for k, v in model.named_parameters(): if not v.requires_grad: continue if k.find("encoder") != -1: net_params.append(v) else: classifier_params.append(v) params = [ { "params": net_params }, { "params": classifier_params, "lr": optimizer_config["classifier_lr"] }, ] else: if master_params: params = master_params else: params = model.parameters() if optimizer_config["type"] == "SGD": optimizer = optim.SGD(params, lr=optimizer_config["learning_rate"], momentum=optimizer_config["momentum"], weight_decay=optimizer_config["weight_decay"], nesterov=optimizer_config["nesterov"]) elif optimizer_config["type"] == "FusedSGD": optimizer = FusedSGD(params, lr=optimizer_config["learning_rate"], momentum=optimizer_config["momentum"], weight_decay=optimizer_config["weight_decay"], nesterov=optimizer_config["nesterov"]) elif optimizer_config["type"] == "Adam": optimizer = optim.Adam(params, lr=optimizer_config["learning_rate"], weight_decay=optimizer_config["weight_decay"]) elif optimizer_config["type"] == "FusedAdam": optimizer = FusedAdam(params, lr=optimizer_config["learning_rate"], weight_decay=optimizer_config["weight_decay"]) elif optimizer_config["type"] == "AdamW": optimizer = AdamW(params, lr=optimizer_config["learning_rate"], weight_decay=optimizer_config["weight_decay"]) elif optimizer_config["type"] == "RmsProp": optimizer = RMSprop(params, lr=optimizer_config["learning_rate"], weight_decay=optimizer_config["weight_decay"]) else: raise KeyError("unrecognized optimizer {}".format( optimizer_config["type"])) if optimizer_config["schedule"]["type"] == "step": scheduler = LRStepScheduler(optimizer, **optimizer_config["schedule"]["params"]) elif optimizer_config["schedule"]["type"] == "clr": scheduler = CyclicLR(optimizer, **optimizer_config["schedule"]["params"]) elif optimizer_config["schedule"]["type"] == "multistep": scheduler = MultiStepLR(optimizer, **optimizer_config["schedule"]["params"]) elif optimizer_config["schedule"]["type"] == "exponential": scheduler = ExponentialLRScheduler( optimizer, **optimizer_config["schedule"]["params"]) elif optimizer_config["schedule"]["type"] == "poly": scheduler = PolyLR(optimizer, **optimizer_config["schedule"]["params"]) elif optimizer_config["schedule"]["type"] == "constant": scheduler = lr_scheduler.LambdaLR(optimizer, lambda epoch: 1.0) elif optimizer_config["schedule"]["type"] == "linear": def linear_lr(it): return it * optimizer_config["schedule"]["params"][ "alpha"] + optimizer_config["schedule"]["params"]["beta"] scheduler = lr_scheduler.LambdaLR(optimizer, linear_lr) return optimizer, scheduler
# PATH RUN_NAME = f"MUDE_{NOISE_TYPE}" TENSORBOARD_RUN_LOG_DIR_PATH = TENSOBOARD_LOGS_DIR_PATH.joinpath(RUN_NAME) CHECKPOINTS_RUN_DIR_PATH = CHECKPOINTS_DIR_PATH.joinpath(RUN_NAME) mude = MUDE(dim=DIM, characters_vocab_size=CHAR_VOCAB_SIZE, tokens_vocab_size=TGT_VOCAB_SIZE, encoder_depth=DEPTH, encoder_attn_heads=ATTN_HEADS, encoder_dimff=DIM_FFT, encoder_dropout=DROPOUT_RATE, top_rec_hidden_dim=DIM_HIDDEN, top_rec_proj_dropout=DROPOUT_RATE) mude = mude.to(device) opt = RMSprop(mude.parameters(), lr=LR) # losses and metric definition seq2seq_criterion = NLLLoss( ignore_index=train_ld.dataset.vect.PAD_CHAR_INDX) recog_criterion = NLLLoss(ignore_index=train_ld.dataset.vocab['<PAD>']) mtrcs = { "WRA": WordRecognitionAccuracy() } # word recognition accuracy, i guess? :) trainer = create_trainer(model=mude, optimizer=opt, seq2seq_loss_fn=seq2seq_criterion, pred_loss_fn=recog_criterion, char_vocab_size=CHAR_VOCAB_SIZE, tgt_vocab_size=TGT_VOCAB_SIZE)
def train(): """ train model """ dataset = dset.ImageFolder(root=opt.dataroot, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, num_workers=int(opt.workers)) if torch.cuda.device_count() > 1: netG = torch.nn.DataParallel(Generator(ngpu)) netD = torch.nn.DataParallel(Discriminator(ngpu)) else: netG = Generator(ngpu) netD = Discriminator(ngpu) netD.apply(weights_init) netG.apply(weights_init) netD.to(device) netG.to(device) if opt.netG != "": netG.load_state_dict( torch.load(opt.netG, map_location=lambda storage, loc: storage)) if opt.netD != "": netD.load_state_dict( torch.load(opt.netD, map_location=lambda storage, loc: storage)) print(netG) print(netD) optimizerD = RMSprop(netD.parameters(), lr=opt.lr) optimizerG = RMSprop(netG.parameters(), lr=opt.lr) fixed_noise = torch.randn(opt.batch_size, opt.nz, 1, 1, device=device) print("########################################") print(f"Train dataset path: {opt.dataroot}") print(f"Batch size: {opt.batch_size}") print(f"Image size: {opt.img_size}") print(f"Epochs: {opt.epochs}") print("########################################") print("Starting trainning!") for epoch in range(opt.epochs): for i, data in enumerate(dataloader): # get data real_imgs = data[0].to(device) batch_size = real_imgs.size(0) # Sample noise as generator input noise = torch.randn(batch_size, 100, 1, 1, device=device) ############################################## # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ############################################## optimizerD.zero_grad() # Generate a batch of images fake_imgs = netG(noise).detach() # Adversarial loss errD = -torch.mean(netD(real_imgs)) + torch.mean(netD(fake_imgs)) errD.backward() optimizerD.step() # Clip weights of discriminator for p in netD.parameters(): p.data.clamp_(-opt.clip_value, opt.clip_value) ############################################## # (2) Update G network: maximize log(D(G(z))) ############################################## if i % opt.n_critic == 0: optimizerG.zero_grad() # Generate a batch of images fake_imgs = netG(noise) # Adversarial loss errG = -torch.mean(netD(fake_imgs)) errG.backward() optimizerG.step() print( f"Epoch->[{epoch + 1:3d}/{opt.epochs}] " f"Progress->[{i}/{len(dataloader)}] " f"Loss_D: {errD.item():.4f} " f"Loss_G: {errG.item():.4f} ", end="\r") if i % 100 == 0: vutils.save_image(real_imgs, f"{opt.outf}/simpson_real_samples.png", normalize=True) fake = netG(fixed_noise) vutils.save_image( fake.detach(), f"{opt.outf}/simpson_fake_samples_epoch_{epoch + 1}.png", normalize=True) # do checkpointing torch.save(netG.state_dict(), f"{opt.checkpoint_dir}/simpson_G.pth") torch.save(netD.state_dict(), f"{opt.checkpoint_dir}/simpson_D.pth")
def train(): """ train model """ try: os.makedirs(opt.checkpoints_dir) except OSError: pass ################################################ # load train dataset ################################################ dataset = dset.CIFAR10(root=opt.dataroot, download=True, transform=transforms.Compose([ transforms.Resize(opt.image_size), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ])) assert dataset dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=True, num_workers=int(opt.workers)) ################################################ # load model ################################################ if torch.cuda.device_count() > 1: netG = torch.nn.DataParallel(Generator(ngpu)) else: netG = Generator(ngpu) if opt.netG != '': netG.load_state_dict(torch.load(opt.netG, map_location=lambda storage, loc: storage)) if torch.cuda.device_count() > 1: netD = torch.nn.DataParallel(Discriminator(ngpu)) else: netD = Discriminator(ngpu) if opt.netD != '': netD.load_state_dict(torch.load(opt.netD, map_location=lambda storage, loc: storage)) # set train mode netG.train() netG = netG.to(device) netD.train() netD = netD.to(device) print(netG) print(netD) ################################################ # Use RMSprop optimizer ################################################ optimizerD = RMSprop(netD.parameters(), lr=opt.lr) optimizerG = RMSprop(netG.parameters(), lr=opt.lr) ################################################ # print args ################################################ print("########################################") print(f"train dataset path: {opt.dataroot}") print(f"work thread: {opt.workers}") print(f"batch size: {opt.batch_size}") print(f"image size: {opt.image_size}") print(f"Epochs: {opt.n_epochs}") print(f"Noise size: {opt.nz}") print("########################################") print("Starting trainning!") for epoch in range(opt.n_epochs): for i, data in enumerate(dataloader): # get data real_imgs = data[0].to(device) batch_size = real_imgs.size(0) # Sample noise as generator input noise = torch.randn(batch_size, nz, 1, 1, device=device) ############################################## # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))) ############################################## optimizerD.zero_grad() # Generate a batch of images fake_imgs = netG(noise).detach() # Adversarial loss real_output = netD(real_imgs) fake_output = netD(fake_imgs) loss_D = -torch.mean(real_output) + torch.mean(fake_output) loss_D.backward() optimizerD.step() # Clip weights of discriminator for p in netD.parameters(): p.data.clamp_(-opt.clip_value, opt.clip_value) ############################################## # (2) Update G network: maximize log(D(G(z))) ############################################## if i % opt.n_critic == 0: optimizerG.zero_grad() # Generate a batch of images fake_imgs = netG(noise) # Adversarial loss loss_G = -torch.mean(netD(fake_imgs)) loss_G.backward() optimizerG.step() print(f"Epoch->[{epoch + 1:03d}/{opt.n_epochs:03d}] " f"Progress->{i / len(dataloader) * 100:4.2f}% " f"Loss_D: {loss_D.item():.4f} " f"Loss_G: {loss_G.item():.4f} ", end="\r") if i % 100 == 0: vutils.save_image(real_imgs, f"{opt.out_images}/real_samples.png", normalize=True) with torch.no_grad(): fake = netG(fixed_noise).detach().cpu() vutils.save_image(fake, f"{opt.out_images}/fake_samples_epoch_{epoch + 1:03d}.png", normalize=True) # do checkpointing torch.save(netG.state_dict(), f"{opt.checkpoints_dir}/netG_epoch_{epoch + 1:03d}.pth") torch.save(netD.state_dict(), f"{opt.checkpoints_dir}/netD_epoch_{epoch + 1:03d}.pth")
def __init__(self, algo_params, env, transition_tuple=None, path=None, seed=-1): # environment self.env = env self.env.seed(seed) obs = self.env.reset() self.frame_skip = algo_params['frame_skip'] self.original_image_shape = obs.shape self.image_size = algo_params['image_size'] algo_params.update({ 'state_shape': (self.frame_skip, self.image_size, self.image_size), 'action_dim': self.env.action_space.n, 'init_input_means': None, 'init_input_vars': None }) # training args self.training_epoch = algo_params['training_epoch'] self.training_frame_per_epoch = algo_params['training_frame_per_epoch'] self.printing_gap = algo_params['printing_gap'] self.testing_gap = algo_params['testing_gap'] self.testing_frame_per_epoch = algo_params['testing_frame_per_epoch'] self.saving_gap = algo_params['saving_gap'] # args for compatibility and are NOT to be used algo_params['actor_learning_rate'] = 0.0 algo_params['observation_normalization'] = False algo_params['tau'] = 1.0 super(DQN, self).__init__(algo_params, transition_tuple=transition_tuple, image_obs=True, action_type='discrete', path=path, seed=seed) # torch self.network_dict.update({ 'Q': DQNetwork(self.state_shape, self.action_dim).to(self.device), 'Q_target': DQNetwork(self.state_shape, self.action_dim).to(self.device) }) self.network_keys_to_save = ['Q', 'Q_target'] self.Q_optimizer = RMSprop(self.network_dict['Q'].parameters(), lr=self.critic_learning_rate, eps=algo_params['RMSprop_epsilon'], weight_decay=algo_params['Q_weight_decay'], centered=True) self._soft_update(self.network_dict['Q'], self.network_dict['Q_target'], tau=1) # behavioural policy args (exploration) epsilong_decay_frame = algo_params[ 'epsilon_decay_fraction'] * self.training_epoch * self.training_frame_per_epoch self.exploration_strategy = LinearDecayGreedy( decay=epsilong_decay_frame, rng=self.rng) # training args self.warmup_step = algo_params['warmup_step'] self.Q_target_update_interval = algo_params['Q_target_update_interval'] self.last_frame = None self.frame_buffer = [None, None, None, None] self.frame_count = 0 self.reward_clip = algo_params['reward_clip'] # statistic dict self.statistic_dict.update({ 'epoch_return': [], 'epoch_test_return': [] })