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 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, 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 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(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 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 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):
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):
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 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_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 __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 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 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_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 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)
# 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 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'))
class Train:
def __init__(self, model, trainloader, valloader, args):
self.model = model
self.trainloader = trainloader
self.valloader = valloader
self.args = args
self.start_epoch = 0
self.best_top1 = 0.0
# Loss function and Optimizer
self.loss = None
self.optimizer = None
self.create_optimization()
# Model Loading
self.load_pretrained_model()
self.load_checkpoint(self.args.resume_from)
# Tensorboard Writer
self.summary_writer = SummaryWriter(log_dir=args.summary_dir)
def train(self):
for cur_epoch in range(self.start_epoch, self.args.num_epochs):
# Initialize tqdm
tqdm_batch = tqdm(self.trainloader,
desc="Epoch-" + str(cur_epoch) + "-")
# Learning rate adjustment
self.adjust_learning_rate(self.optimizer, cur_epoch)
# Meters for tracking the average values
loss, top1, top5 = AverageTracker(), AverageTracker(
), AverageTracker()
# Set the model to be in training mode (for dropout and batchnorm)
self.model.train()
for data, target in tqdm_batch:
if self.args.cuda:
data, target = data.cuda(), target.cuda()
data_var, target_var = Variable(data), Variable(target)
# Forward pass
output = self.model(data_var)
cur_loss = self.loss(output, target_var)
# Optimization step
self.optimizer.zero_grad()
cur_loss.backward()
self.optimizer.step()
# Top-1 and Top-5 Accuracy Calculation
cur_acc1, cur_acc5 = self.compute_accuracy(output.data,
target,
topk=(1, 5))
loss.update(cur_loss.data[0])
top1.update(cur_acc1[0])
top5.update(cur_acc5[0])
# Summary Writing
self.summary_writer.add_scalar("epoch-loss", loss.avg, cur_epoch)
self.summary_writer.add_scalar("epoch-top-1-acc", top1.avg,
cur_epoch)
self.summary_writer.add_scalar("epoch-top-5-acc", top5.avg,
cur_epoch)
# Print in console
tqdm_batch.close()
print("Epoch-" + str(cur_epoch) + " | " + "loss: " +
str(loss.avg) + " - acc-top1: " + str(top1.avg)[:7] +
"- acc-top5: " + str(top5.avg)[:7])
# Evaluate on Validation Set
if cur_epoch % self.args.test_every == 0 and self.valloader:
self.test(self.valloader, cur_epoch)
# Checkpointing
is_best = top1.avg > self.best_top1
self.best_top1 = max(top1.avg, self.best_top1)
self.save_checkpoint(
{
'epoch': cur_epoch + 1,
'state_dict': self.model.state_dict(),
'best_top1': self.best_top1,
'optimizer': self.optimizer.state_dict(),
}, is_best)
def test(self, testloader, cur_epoch=-1):
loss, top1, top5 = AverageTracker(), AverageTracker(), AverageTracker()
# Set the model to be in testing mode (for dropout and batchnorm)
self.model.eval()
for data, target in testloader:
if self.args.cuda:
data, target = data.cuda(), target.cuda()
data_var, target_var = Variable(data, volatile=True), Variable(
target, volatile=True)
# Forward pass
output = self.model(data_var)
cur_loss = self.loss(output, target_var)
# Top-1 and Top-5 Accuracy Calculation
cur_acc1, cur_acc5 = self.compute_accuracy(output.data,
target,
topk=(1, 5))
loss.update(cur_loss.data[0])
top1.update(cur_acc1[0])
top5.update(cur_acc5[0])
if cur_epoch != -1:
# Summary Writing
self.summary_writer.add_scalar("test-loss", loss.avg, cur_epoch)
self.summary_writer.add_scalar("test-top-1-acc", top1.avg,
cur_epoch)
self.summary_writer.add_scalar("test-top-5-acc", top5.avg,
cur_epoch)
print("Test Results" + " | " + "loss: " + str(loss.avg) +
" - acc-top1: " + str(top1.avg)[:7] + "- acc-top5: " +
str(top5.avg)[:7])
def save_checkpoint(self, state, is_best, filename='checkpoint.pth.tar'):
torch.save(state, self.args.checkpoint_dir + filename)
if is_best:
shutil.copyfile(self.args.checkpoint_dir + filename,
self.args.checkpoint_dir + 'model_best.pth.tar')
def compute_accuracy(self, output, target, topk=(1, )):
"""Computes the accuracy@k for the specified values of k"""
maxk = max(topk)
batch_size = target.size(0)
_, idx = output.topk(maxk, 1, True, True)
idx = idx.t()
correct = idx.eq(target.view(1, -1).expand_as(idx))
acc_arr = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
acc_arr.append(correct_k.mul_(1.0 / batch_size))
return acc_arr
def adjust_learning_rate(self, optimizer, epoch):
"""Sets the learning rate to the initial LR multiplied by 0.98 every epoch"""
learning_rate = self.args.learning_rate * (
self.args.learning_rate_decay**epoch)
for param_group in optimizer.param_groups:
param_group['lr'] = learning_rate
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 load_pretrained_model(self):
try:
print("Loading ImageNet pretrained weights...")
pretrained_dict = torch.load(self.args.pretrained_path)
self.model.load_state_dict(pretrained_dict)
print("ImageNet pretrained weights loaded successfully.\n")
except:
print("No ImageNet pretrained weights exist. Skipping...\n")
def load_checkpoint(self, filename):
filename = self.args.checkpoint_dir + filename
try:
print("Loading checkpoint '{}'".format(filename))
checkpoint = torch.load(filename)
self.start_epoch = checkpoint['epoch']
self.best_top1 = checkpoint['best_top1']
self.model.load_state_dict(checkpoint['state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer'])
print("Checkpoint loaded successfully from '{}' at (epoch {})\n".
format(self.args.checkpoint_dir, checkpoint['epoch']))
except:
print("No checkpoint exists from '{}'. Skipping...\n".format(
self.args.checkpoint_dir))
class DQN(Agent):
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': []
})
def run(self, test=False, render=False, load_network_ep=None, sleep=0):
if test:
num_frames = self.testing_frame_per_epoch
if load_network_ep is not None:
print("Loading network parameters...")
self._load_network(ep=load_network_ep)
print("Start testing...")
else:
num_frames = self.training_frame_per_epoch
print("Start training...")
for epo in range(self.training_epoch):
ep_return = self._interact(render,
test,
epo=epo,
num_frames=num_frames,
sleep=sleep)
self.statistic_dict['epoch_return'].append(ep_return)
print("Finished training epoch %i, " % epo,
"full return %0.1f" % ep_return)
if (epo % self.testing_gap == 0) and (epo != 0) and (not test):
print("Evaluate agent at epoch %i..." % epo)
ep_test_return = self._interact(
render,
test=True,
epo=epo,
num_frames=self.testing_frame_per_epoch)
self.statistic_dict['epoch_test_return'].append(ep_test_return)
print("Finished testing epoch %i, " % epo,
"test return %0.1f" % ep_test_return)
if (epo % self.saving_gap == 0) and (epo != 0) and (not test):
self._save_network(ep=epo)
if not test:
print("Finished training")
print("Saving statistics...")
self._save_statistics()
self._plot_statistics()
else:
print("Finished testing")
def _interact(self,
render=False,
test=False,
epo=0,
num_frames=0,
sleep=0):
ep_return = 0
self.frame_count = 0
while self.frame_count < num_frames:
done = False
obs = self.env.reset()
obs = self._pre_process([obs])
num_lives = self.env.ale.lives()
# start a new episode
while not done:
if render:
self.env.render()
if self.env_step_count < self.warmup_step:
action = self.env.action_space.sample()
else:
action = self._select_action(obs, test=test)
# action repeat, aggregated reward
frames = []
added_reward = 0
for _ in range(self.frame_skip):
new_obs, reward, done, info = self.env.step(action)
frames.append(new_obs.copy())
added_reward += reward
time.sleep(sleep)
# frame gray scale, resize, stack
new_obs = self._pre_process(frames[-2:])
# reward clipped into [-1, 1]
reward = max(min(added_reward, self.reward_clip),
-self.reward_clip)
if num_lives > self.env.ale.lives():
# treat the episode as terminated when the agent loses a live in the game
num_lives = self.env.ale.lives()
done_to_save = True
# set the reward to be -reward_bound
reward = -self.reward_clip
# clear frame buffer when the agent starts with a new live
self.frame_buffer = [None, None, None, None]
else:
done_to_save = done
# return to be recorded
ep_return += reward
if not test:
self._remember(obs, action, new_obs, reward,
1 - int(done_to_save))
if (self.env_step_count % self.update_interval
== 0) and (self.env_step_count > self.warmup_step):
self._learn()
obs = new_obs
self.frame_count += 1
self.env_step_count += 1
if self.frame_count % self.printing_gap == 0 and self.frame_count != 0:
print("Epoch %i" % epo,
"passed frames %i" % self.frame_count,
"return %0.1f" % ep_return)
# clear frame buffer at the end of an episode
self.frame_buffer = [None, None, None, None]
return ep_return
def _select_action(self, obs, test=False):
if test:
obs = T.tensor([obs], dtype=T.float32).to(self.device)
with T.no_grad():
action = self.network_dict['Q_target'].get_action(obs)
return action
else:
if self.exploration_strategy(self.env_step_count):
action = self.rng.integers(self.action_dim)
else:
obs = T.tensor([obs], dtype=T.float32).to(self.device)
with T.no_grad():
action = self.network_dict['Q_target'].get_action(obs)
return action
def _learn(self, steps=None):
if len(self.buffer) < self.batch_size:
return
if steps is None:
steps = self.optimizer_steps
for i in range(steps):
if self.prioritised:
batch, weights, inds = self.buffer.sample(self.batch_size)
weights = T.tensor(weights).view(self.batch_size,
1).to(self.device)
else:
batch = self.buffer.sample(self.batch_size)
weights = T.ones(size=(self.batch_size, 1)).to(self.device)
inds = None
inputs = T.tensor(batch.state, dtype=T.float32).to(self.device)
actions = T.tensor(batch.action,
dtype=T.long).unsqueeze(1).to(self.device)
inputs_ = T.tensor(batch.next_state,
dtype=T.float32).to(self.device)
rewards = T.tensor(batch.reward,
dtype=T.float32).unsqueeze(1).to(self.device)
done = T.tensor(batch.done,
dtype=T.float32).unsqueeze(1).to(self.device)
if self.discard_time_limit:
done = done * 0 + 1
with T.no_grad():
maximal_next_values = self.network_dict['Q_target'](
inputs_).max(1)[0].view(self.batch_size, 1)
value_target = rewards + done * self.gamma * maximal_next_values
self.Q_optimizer.zero_grad()
value_estimate = self.network_dict['Q'](inputs).gather(1, actions)
loss = F.smooth_l1_loss(value_estimate,
value_target.detach(),
reduction='none')
(loss * weights).mean().backward()
self.Q_optimizer.step()
if self.prioritised:
assert inds is not None
self.buffer.update_priority(
inds, np.abs(loss.cpu().detach().numpy()))
self.statistic_dict['critic_loss'].append(
loss.detach().mean().cpu().numpy().item())
if self.optim_step_count % self.Q_target_update_interval == 0:
self._soft_update(self.network_dict['Q'],
self.network_dict['Q_target'],
tau=1)
self.optim_step_count += 1
def _pre_process(self, frames):
# This method takes 2 frames and does the following things
# 1. Max-pool two consecutive frames to deal with flickering
# 2. Convert images to Y channel: Y = 0.299*R + 0.587*G + (1 - (0.299 + 0.587))*B
# 3. Resize images to 84x84
# 4. Stack it with previous frames as one observation
# output: 1000, 1200, 1230, 1234, 2345, 3456...
if len(frames) == 1:
frames.insert(0, np.zeros(self.original_image_shape))
assert len(frames) == 2
last_img = frames[0].copy()
img = frames[1].copy()
img = np.max([last_img, img], axis=0)
img = img.transpose((-1, 0, 1))
img_Y = 0.299 * img[0] + 0.587 * img[1] + (1 -
(0.299 + 0.587)) * img[2]
img_Y_resized = np.asarray(
Image.fromarray(img_Y).resize((self.image_size, self.image_size),
Image.BILINEAR))
for i in range(len(self.frame_buffer)):
if self.frame_buffer[i] is None:
self.frame_buffer[i] = img_Y_resized.copy()
break
if i == (len(self.frame_buffer) - 1):
del self.frame_buffer[0]
self.frame_buffer.append(img_Y_resized.copy())
obs = []
for i in range(len(self.frame_buffer)):
if self.frame_buffer[i] is not None:
obs.append(self.frame_buffer[i].copy())
else:
obs.append(np.zeros((self.image_size, self.image_size)))
return np.array(obs, dtype=np.uint8)
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 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
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': []
})
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 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")
