def init_evaluation_container(self):
# Define network
self.model = generate_net(self.args.models)
# # Resuming checkpoint
state_dict, _, _, _ = load_checkpoint(checkpoint_path=self.args.evaluation.resume_eval)
load_pretrained_model(self.model, state_dict)
self.model = self.model.cuda()
def init_training_container(self):
# Define dataset
self.train_set = self.Dataset_train(self.args.dataset, split='train')
self.val_set = self.Dataset_val(self.args.dataset, split='val')
self.test_set = self.Dataset_val(self.args.dataset, split='test')
# Define network
self.model = generate_net(self.args.models)
self.model = self.model.cuda()
start_it = 0
stage = 0
# # Resuming checkpoint
if self.args.training.resume_train is not None and os.path.exists(self.args.training.resume_train):
state_dict, optimizer, start_it, stage = load_checkpoint(checkpoint_path=self.args.training.resume_train)
self.gen_optimizer(self.model.param_groups(), (stage + 1) // 2)
if isinstance(self.args.training.batchsize, list):
self.batchsize = self.args.training.batchsize[(stage + 1) // 2]
else:
self.batchsize = self.args.training.batchsize
load_pretrained_model(self.model, state_dict)
self.model = self.model.cuda()
if not self.args.training.ft and optimizer is not None:
for name in self.args.training.optimizer.keys():
if name in optimizer.keys():
self.optimizer[name].load_state_dict(optimizer[name])
else:
start_it = 0
else:
self.gen_optimizer(self.model.param_groups(), (stage + 1) // 2)
self.start_epoch = 0
self.start_it = start_it
self.stage = stage
# Define Criterion
self.criterion = MMLoss(self.args.training)
def main():
global args
print("config:{0}".format(args))
checkpoint_dir = args.checkpoint_folder
global_step = 0
min_val_loss = 999999999
title = 'train|val loss '
init = np.NaN
win_feats5 = viz.line(
X=np.column_stack((np.array([init]), np.array([init]))),
Y=np.column_stack((np.array([init]), np.array([init]))),
opts={
'title': title,
'xlabel': 'Iter',
'ylabel': 'Loss',
'legend': ['train_feats5', 'val_feats5']
},
)
win_fusion = viz.line(
X=np.column_stack((np.array([init]), np.array([init]))),
Y=np.column_stack((np.array([init]), np.array([init]))),
opts={
'title': title,
'xlabel': 'Iter',
'ylabel': 'Loss',
'legend': ['train_fusion', 'val_fusion']
},
)
train_loader, val_loader = prep_SBD_dataset.get_dataloader(args)
model = CASENet_resnet101(pretrained=False, num_classes=args.cls_num)
if args.multigpu:
model = torch.nn.DataParallel(model.cuda())
else:
model = model.cuda()
policies = get_model_policy(model) # Set the lr_mult=10 of new layer
optimizer = torch.optim.SGD(policies,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
cudnn.benchmark = True
if args.pretrained_model:
utils.load_pretrained_model(model, args.pretrained_model)
if args.resume_model:
checkpoint = torch.load(args.resume_model)
args.start_epoch = checkpoint['epoch'] + 1
min_val_loss = checkpoint['min_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
for epoch in range(args.start_epoch, args.epochs):
curr_lr = utils.adjust_learning_rate(args.lr, args, optimizer,
global_step, args.lr_steps)
global_step = model_play.train(args, train_loader, model, optimizer, epoch, curr_lr,\
win_feats5, win_fusion, viz, global_step)
curr_loss = model_play.validate(args, val_loader, model, epoch,
win_feats5, win_fusion, viz,
global_step)
# Always store current model to avoid process crashed by accident.
utils.save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'min_loss': min_val_loss,
},
epoch,
folder=checkpoint_dir,
filename="curr_checkpoint.pth.tar")
if curr_loss < min_val_loss:
min_val_loss = curr_loss
utils.save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'min_loss': min_val_loss,
},
epoch,
folder=checkpoint_dir)
print("Min loss is {0}, in {1} epoch.".format(min_val_loss, epoch))
]
else:
image_file = os.path.join(args.image_dir, args.image_file)
if os.path.exists(image_file):
ori_test_list = [args.image_file]
test_list = [image_file]
else:
raise IOError('nothing to be tested!')
# load network
num_cls = 19
model = CASENet_resnet101(pretrained=False, num_classes=num_cls)
# model = model.cuda()
model = model.eval()
# cudnn.benchmark = True
utils.load_pretrained_model(model, args.model)
cls_names = get_cityscapes_class_names()
color_dict = get_colors()
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
# Define normalization for data
input_size = 472
normalize = transforms.Normalize(mean=[104.008, 116.669, 122.675],
std=[1, 1, 1])
img_transform = transforms.Compose([
transforms.Resize([input_size, input_size]),
RGB2BGR(roll=True),
def main():
global args
print("config:{0}".format(args))
torch.manual_seed(args.seed)
if torch.cuda.is_available():
if not args.cuda:
print("WARNING: You have a CUDA device, so you should probably run with --cuda")
device = torch.device("cuda" if args.cuda else "cpu")
global_step = 0
min_loss = 999999
min_loss_epoch = 0
model_writer = SummaryWriter(log_dir=args.log_dir)
# train_loader, val_loader = prep_discrete_ori_audio_pose_data.get_dataloader(args)
train_loader, val_loader = prep_audio_discrete_data.get_dataloader(args)
# Define Model
model = TransformerDiscreteDecoder(n_dec_layers=args.n_dec_layers, n_head=args.n_head,
d_feats=args.feats_dim, d_model=args.d_model, d_k=args.d_k,
d_v=args.d_v, d_out=args.d_out, num_cls=args.num_cls, max_timesteps=args.max_timesteps,
temperature=args.temperature, add_mfcc=args.add_mfcc, add_beat=args.add_beat, multi_stream=args.multi_stream).to(device)
if args.test_model and args.vis_results:
utils.load_pretrained_model(model, args.test_model)
if args.add_mfcc or args.add_beat:
two_stream_transformer_discrete_model_play.add_audio_given_start_inference_vis(args, val_loader, model, \
device, -1, model_writer, global_step)
else:
two_stream_transformer_discrete_model_play.given_start_inference_vis(args, val_loader, model, \
device, -1, model_writer, global_step)
return
# Define Optimizer
optimizer = torch.optim.Adam(list(model.parameters()), lr=args.lr, weight_decay=1e-5)
ce_criterion = maskedCrossEntropy
curr_exist_model = os.path.join(args.checkpoint_folder, "curr_checkpoint.pth.tar")
if os.path.exists(curr_exist_model):
checkpoint = torch.load(curr_exist_model)
args.start_epoch = checkpoint['epoch']+1
min_loss = checkpoint['min_loss']
utils.load_pretrained_model(model, curr_exist_model)
optimizer.load_state_dict(checkpoint['optimizer'])
curr_loss = two_stream_transformer_discrete_model_play.validate(args, val_loader, model, ce_criterion, \
device, -1, model_writer, global_step)
for epoch in range(args.start_epoch, args.epochs):
curr_lr = utils.adjust_learning_rate(args.lr, args, optimizer, epoch, args.lr_steps)
global_step, curr_loss = two_stream_transformer_discrete_model_play.train(args, train_loader, model, optimizer, \
ce_criterion, device, epoch, curr_lr, model_writer, global_step)
if not args.debug:
curr_loss = two_stream_transformer_discrete_model_play.validate(args, val_loader, model, ce_criterion, \
device, epoch, model_writer, global_step)
if curr_loss < min_loss:
min_loss = curr_loss
min_loss_epoch = epoch
utils.save_checkpoint({
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'min_loss': min_loss,
}, epoch, min_loss, folder=args.checkpoint_folder)
print("Min Loss is {0}, in {1} epoch.".format(min_loss, epoch))
# Always save current one.
utils.save_checkpoint({
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'min_loss': min_loss,
}, min_loss_epoch, min_loss, folder=args.checkpoint_folder, filename='curr_checkpoint.pth.tar')
model_writer.close()