def __init__(self, config=None, train_loader=None, test_dataset=None):
self.config = config
self.train_loader = train_loader
self.test_dataset = test_dataset
# model
self.model = FCSN(self.config.n_class)
# optimizer
if self.config.mode == 'train':
self.optimizer = optim.Adam(self.model.parameters())
# self.optimizer = optim.SGD(self.model.parameters(), lr=config.lr, momentum=self.config.momentum)
self.model.train()
if self.config.gpu:
self.model = self.model.cuda()
if not os.path.exists(self.config.score_dir):
os.mkdir(self.config.score_dir)
if not os.path.exists(self.config.save_dir):
os.mkdir(self.config.save_dir)
if not os.path.exists(self.config.log_dir):
os.mkdir(self.config.log_dir)
def __init__(self, config=None, train_loader=None, test_loader=None):
self.config = config
self.train_loader = train_loader
self.test_loader = test_loader
# model
self.model = FCSN(self.config.n_class)
# optimizer
if self.config.mode == 'train':
self.optimizer = optim.Adam(self.model.parameters())
# self.optimizer = optim.SGD(self.model.parameters(), lr=config.lr, momentum=0.9)
self.model.train()
# weight
self.tvsum_weight = torch.tensor([0.55989996, 4.67362574])
if self.config.gpu:
self.model = self.model.cuda()
self.tvsum_weight = self.tvsum_weight.cuda()
class Solver(object):
def __init__(self, config=None, train_loader=None, test_dataset=None):
self.config = config
self.train_loader = train_loader
self.test_dataset = test_dataset
# model
self.model = FCSN(self.config.n_class)
# optimizer
if self.config.mode == 'train':
self.optimizer = optim.Adam(self.model.parameters())
# self.optimizer = optim.SGD(self.model.parameters(), lr=config.lr, momentum=self.config.momentum)
self.model.train()
if self.config.gpu:
self.model = self.model.cuda()
if not os.path.exists(self.config.score_dir):
os.mkdir(self.config.score_dir)
if not os.path.exists(self.config.save_dir):
os.mkdir(self.config.save_dir)
if not os.path.exists(self.config.log_dir):
os.mkdir(self.config.log_dir)
@staticmethod
def sum_loss(pred_score, gt_labels, weight=None):
n_batch, n_class, n_frame = pred_score.shape
log_p = torch.log_softmax(pred_score, dim=1).reshape(-1, n_class)
gt_labels = gt_labels.reshape(-1)
criterion = torch.nn.NLLLoss(weight)
loss = criterion(log_p, gt_labels)
return loss
def train(self):
writer = SummaryWriter(log_dir=self.config.log_dir)
t = trange(self.config.n_epochs, desc='Epoch', ncols=80)
for epoch_i in t:
sum_loss_history = []
for batch_i, (feature, label, _) in enumerate(
tqdm(self.train_loader,
desc='Batch',
ncols=80,
leave=False)):
# [batch_size, 1024, seq_len]
feature.requires_grad_()
# => cuda
if self.config.gpu:
feature = feature.cuda()
label = label.cuda()
# ---- Train ---- #
pred_score = self.model(feature)
label_1 = torch.true_divide(label.sum(), label.shape[0])
label_0 = label.shape[1] - label_1
weight = torch.tensor([label_1, label_0], dtype=torch.float)
if self.config.gpu:
weight = weight.cuda()
loss = self.sum_loss(pred_score, label, weight)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
sum_loss_history.append(loss)
mean_loss = torch.stack(sum_loss_history).mean().item()
t.set_postfix(loss=mean_loss)
writer.add_scalar('Loss', mean_loss, epoch_i)
if (epoch_i + 1) % 5 == 0:
ckpt_path = self.config.save_dir + '/epoch-{}.pkl'.format(
epoch_i)
tqdm.write('Save parameters at {}'.format(ckpt_path))
torch.save(self.model.state_dict(), ckpt_path)
self.evaluate(epoch_i)
self.model.train()
def evaluate(self, epoch_i):
self.model.eval()
out_dict = {}
eval_arr = []
table = PrettyTable()
table.title = 'Eval result of epoch {}'.format(epoch_i)
table.field_names = ['ID', 'Precision', 'Recall', 'F-score']
table.float_format = '1.3'
with h5py.File(self.config.data_path) as data_file:
for feature, label, idx in tqdm(self.test_dataset,
desc='Evaluate',
ncols=80,
leave=False):
if self.config.gpu:
feature = feature.cuda()
pred_score = self.model(feature.unsqueeze(0)).squeeze(0)
pred_score = torch.softmax(pred_score, dim=0)[1]
video_info = data_file['video_' + str(idx)]
pred_score, pred_selected, pred_summary = eval.select_keyshots(
video_info, pred_score)
true_summary_arr = video_info['user_summary'][()]
eval_res = [
eval.eval_metrics(pred_summary, true_summary)
for true_summary in true_summary_arr
]
eval_res = np.mean(eval_res, axis=0).tolist()
eval_arr.append(eval_res)
table.add_row([idx] + eval_res)
out_dict[idx] = {
'pred_score': pred_score,
'pred_selected': pred_selected,
'pred_summary': pred_summary
}
score_save_path = self.config.score_dir + '/epoch-{}.json'.format(
epoch_i)
with open(score_save_path, 'w') as f:
tqdm.write('Save score at {}'.format(str(score_save_path)))
json.dump(out_dict, f)
eval_mean = np.mean(eval_arr, axis=0).tolist()
table.add_row(['mean'] + eval_mean)
tqdm.write(str(table))
class Solver(object):
"""Class that Builds, Trains FCSN model"""
def __init__(self, config=None, train_loader=None, test_loader=None):
self.config = config
self.train_loader = train_loader
self.test_loader = test_loader
# model
self.model = FCSN(self.config.n_class)
# optimizer
if self.config.mode == 'train':
self.optimizer = optim.Adam(self.model.parameters())
# self.optimizer = optim.SGD(self.model.parameters(), lr=config.lr, momentum=0.9)
self.model.train()
# weight
self.tvsum_weight = torch.tensor([0.55989996, 4.67362574])
if self.config.gpu:
self.model = self.model.cuda()
self.tvsum_weight = self.tvsum_weight.cuda()
@staticmethod
def sum_loss(pred_score, gt_labels, weight=None):
n_batch, n_class, n_frame = pred_score.shape
log_p = torch.log_softmax(pred_score, dim=1).reshape(-1, n_class)
gt_labels = gt_labels.reshape(-1)
criterion = torch.nn.NLLLoss(weight)
loss = criterion(log_p, gt_labels)
return loss
def train(self):
writer = SummaryWriter()
for epoch_i in trange(self.config.n_epochs, desc='Epoch', ncols=80):
sum_loss_history = []
for batch_i, (feature, label, _) in enumerate(
tqdm(self.train_loader,
desc='Batch',
ncols=80,
leave=False)):
# [batch_size, 1024, seq_len]
feature.requires_grad_()
# => cuda
if self.config.gpu:
feature = feature.cuda()
label = label.cuda()
# ---- Train ---- #
pred_score = self.model(feature)
# pred_label = torch.argmax(pred_score, dim=1)
# loss = torch.nn.MSELoss(pred_label, label)
loss = self.sum_loss(pred_score, label, self.tvsum_weight)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
sum_loss_history.append(loss)
mean_loss = torch.stack(sum_loss_history).mean().item()
tqdm.write('\nEpoch {}'.format(epoch_i))
tqdm.write('sum loss: {:.3f}'.format(mean_loss))
writer.add_scalar('Loss', mean_loss, epoch_i)
if (epoch_i + 1) % 10 == 0:
if not os.path.exists(self.config.save_dir):
os.mkdir(self.config.save_dir)
ckpt_path = self.config.save_dir + '/epoch-{}.pkl'.format(
epoch_i)
tqdm.write('Save parameters at {}'.format(ckpt_path))
torch.save(self.model.state_dict(), ckpt_path)
self.evaluate(epoch_i)
self.model.train()
def evaluate(self, epoch_i):
self.model.eval()
out_dict = {}
for feature, _, idx in tqdm(self.test_loader,
desc='Evaluate',
ncols=80,
leave=False):
if self.config.gpu:
feature = feature.cuda()
pred_score = self.model(feature.unsqueeze(0))
pred_label = torch.argmax(pred_score,
dim=1).squeeze(0).type(dtype=torch.int)
pred_label = np.array(pred_label.cpu().data).tolist()
out_dict[idx] = pred_label
if not os.path.exists(self.config.score_dir):
os.mkdir(self.config.score_dir)
score_save_path = self.config.score_dir + '/epoch-{}.json'.format(
epoch_i)
with open(score_save_path, 'w') as f:
tqdm.write('Saving score at {}.'.format(str(score_save_path)))
json.dump(out_dict, f)
class Solver(object):
"""Class that Builds, Trains FCSN model"""
def __init__(self, config=None, train_loader=None, test_dataset=None):
self.config = config
self.train_loader = train_loader
self.test_dataset = test_dataset
# model
self.model = FCSN(self.config.n_class)
# optimizer
if self.config.mode == 'train':
self.optimizer = optim.Adam(self.model.parameters())
# self.optimizer = optim.SGD(self.model.parameters(), lr=config.lr, momentum=self.config.momentum)
self.model.train()
if self.config.gpu:
self.model = self.model.cuda()
if not os.path.exists(self.config.score_dir):
os.mkdir(self.config.score_dir)
if not os.path.exists(self.config.save_dir):
os.mkdir(self.config.save_dir)
if not os.path.exists(self.config.log_dir):
os.mkdir(self.config.log_dir)
@staticmethod
def sum_loss(pred_score, gt_labels, weight=None):
n_batch, n_class, n_frame = pred_score.shape
log_p = torch.log_softmax(pred_score, dim=1).reshape(-1, n_class)
gt_labels = gt_labels.reshape(-1)
criterion = torch.nn.NLLLoss(weight)
loss = criterion(log_p, gt_labels)
return loss
def evaluate(self, model_path):
self.model.eval()
out_dict = {}
eval_arr = []
table = PrettyTable()
checkpoint = torch.load(model_path)
self.model.load_state_dict(checkpoint)
table.title = 'Eval result'
table.field_names = ['ID', 'Precision', 'Recall', 'F-score']
table.float_format = '1.3'
inference_time = []
with h5py.File(self.config.data_path) as data_file:
for feature, label, idx in tqdm(self.test_dataset,
desc='Evaluate',
ncols=80,
leave=False):
if self.config.gpu:
feature = feature.cuda()
start = time.time()
#print(feature.size())
pred_score = self.model(feature.unsqueeze(0)).squeeze(0)
inference_time.append(time.time() - start)
pred_score = torch.softmax(pred_score, dim=0)[1]
video_info = data_file['video_' + str(idx)]
pred_score, pred_selected, pred_summary = eval.select_keyshots(
video_info, pred_score)
true_summary_arr = video_info['user_summary'][()]
#print(len(pred_summary), len(true_summary_arr[0]))
eval_res = [
eval.eval_metrics(pred_summary, true_summary)
for true_summary in true_summary_arr
]
eval_res = np.mean(eval_res, axis=0).tolist()
eval_arr.append(eval_res)
table.add_row([idx] + eval_res)
out_dict[idx] = {
'pred_score': pred_score,
'pred_selected': pred_selected,
'pred_summary': pred_summary
}
eval_mean = np.mean(eval_arr, axis=0).tolist()
table.add_row(['mean'] + eval_mean)
tqdm.write(str(table))
print(inference_time)