def __init__(self, config, net=None):
self.log_dir = config.log_dir
self.model_dir = config.model_dir
self.net = net
self.best_val_loss = np.inf
self.tb_writer = SummaryWriter(log_dir=self.log_dir)
self.clock = TrainClock()
def __init__(self, train_spec, net=None):
# setproctitle(config.exp_name)
self.log_dir = train_spec.log_dir
ensure_dir(self.log_dir)
# logconf.set_output_file(os.path.join(self.log_dir, 'log.txt'))
self.model_dir = train_spec.log_model_dir
ensure_dir(self.model_dir)
self.net = net
self.clock = TrainClock()
def __init__(self, config):
self.model_dir = config.model_dir
self.clock = TrainClock()
self.batch_size = config.batch_size
# build network
self.net = self.build_net(config)
# set loss function
self.set_loss_function()
# set optimizer
self.set_optimizer(config)
def __init__(self, config):
self.log_dir = config.log_dir
self.model_dir = config.model_dir
self.clock = TrainClock()
self.device = config.device
self.batch_size = config.batch_size
# build network
self.net = self.build_net()
# set loss function
self.set_loss_function()
# set optimizer
self.set_optimizer(config)
# set tensorboard writer
self.train_tb = SummaryWriter(os.path.join(self.log_dir, 'train.events'))
self.val_tb = SummaryWriter(os.path.join(self.log_dir, 'val.events'))
def __init__(self, config):
self.log_dir = config.log_dir
self.model_dir = config.model_dir
self.clock = TrainClock()
self.batch_size = config.batch_size
# build network
self.net = self.build_net(config)
# print('-----network architecture-----')
# print(self.net)
# set loss function
self.set_loss_function()
# set optimizer and scheduler
self.set_optimizer(config)
self.set_scheduler(config)
# set tensorboard writer
self.train_tb = SummaryWriter(os.path.join(self.log_dir, 'train.events'))
self.val_tb = SummaryWriter(os.path.join(self.log_dir, 'val.events'))
parser = argparse.ArgumentParser()
parser.add_argument('--val', action='store_true', help='weither to run validation')
parser.add_argument('--bench', action='store_true',
help = 'weither to generate results on benchmark dataset')
parser.add_argument('--smooth', action='store_true',
help = 'weither to generator smootGrad results on benchmark dataset')
parser.add_argument('--gen', action = 'store_true',
help = 'weither to generate Large eps adversarial examples on benchmark data')
parser.add_argument('--resume', type=str, default=None,
help='checkpoint path')
args = parser.parse_args()
clock = TrainClock()
clock.epoch = 21
net = ant_model()
net.cuda()
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
check_point = torch.load(args.resume)
net.load_state_dict(check_point['state_dict'])
print('Modeled loaded from {} with metrics:'.format(args.resume))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
base_path = os.path.split(args.resume)[0]
else:
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-c',
'--continue',
dest='continue_path',
type=str,
required=False)
parser.add_argument('-g', '--gpu_ids', type=int, default=0, required=False)
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_ids)
config.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
config.isTrain = True
if not os.path.exists('train_log'):
os.symlink(config.exp_dir, 'train_log')
# get dataset
train_loader = get_dataloader("train", batch_size=config.batch_size)
val_loader = get_dataloader("test", batch_size=config.batch_size)
val_cycle = cycle(val_loader)
dataset_size = len(train_loader)
print('The number of training motions = %d' %
(dataset_size * config.batch_size))
# create tensorboard writer
train_tb = SummaryWriter(os.path.join(config.log_dir, 'train.events'))
val_tb = SummaryWriter(os.path.join(config.log_dir, 'val.events'))
# get model
net = CycleGANModel(config)
net.print_networks(True)
# start training
clock = TrainClock()
net.train()
for e in range(config.nr_epochs):
# begin iteration
pbar = tqdm(train_loader)
for b, data in enumerate(pbar):
net.train()
net.set_input(
data) # unpack data from dataset and apply preprocessing
net.optimize_parameters(
) # calculate loss functions, get gradients, update network weights
# get loss
losses_values = net.get_current_losses()
# update tensorboard
train_tb.add_scalars('train_loss',
losses_values,
global_step=clock.step)
# visualize
if clock.step % config.visualize_frequency == 0:
motion_dict = net.infer()
for k, v in motion_dict.items():
phase = 'h' if k[-1] == 'A' else 'nh'
motion3d = train_loader.dataset.preprocess_inv(
v.detach().cpu().numpy()[0], phase)
img = plot_motion(motion3d, phase)
train_tb.add_image(k, img, global_step=clock.step)
pbar.set_description("EPOCH[{}][{}/{}]".format(
e, b, len(train_loader)))
pbar.set_postfix(OrderedDict(losses_values))
# validation
if clock.step % config.val_frequency == 0:
net.eval()
data = next(val_cycle)
net.set_input(data)
net.forward()
losses_values = net.get_current_losses()
val_tb.add_scalars('val_loss',
losses_values,
global_step=clock.step)
# visualize
if clock.step % config.visualize_frequency == 0:
motion_dict = net.infer()
for k, v in motion_dict.items():
phase = 'h' if k[-1] == 'A' else 'nh'
motion3d = val_loader.dataset.preprocess_inv(
v.detach().cpu().numpy()[0], phase)
img = plot_motion(motion3d, phase)
val_tb.add_image(k, img, global_step=clock.step)
clock.tick()
# leraning_rate to tensorboarrd
lr = net.optimizers[0].param_groups[0]['lr']
train_tb.add_scalar("learning_rate", lr, global_step=clock.step)
if clock.epoch % config.save_frequency == 0:
net.save_networks(epoch=e)
clock.tock()
net.update_learning_rate(
) # update learning rates at the end of every epoch.