def train(**kwargs):
opt._parse(kwargs)
model = models.Key_CSVideoNet(opt.CR[0],opt.Height,opt.Width,1)
#if t.cuda.device_count() > 1:
# model = t.nn.DataParallel(model,device_ids = [0,1,2,3])
model.to(opt.device)
train_data = UCF101(opt.train_data_root,train=True)
val_data = UCF101(opt.train_data_root,train=False)
train_dataloader = DataLoader(train_data,opt.batch_size,
shuffle=True,num_workers=opt.num_workers)
val_dataloader = DataLoader(val_data,opt.batch_size,
shuffle=False,num_workers=opt.num_workers)
criterion = t.nn.MSELoss()
lr = opt.lr
optimizer = t.optim.Adam(model.parameters(),
lr = lr,
weight_decay = opt.weight_decay)
utils_eval = utils.Evaluation()
loss_meter = meter.AverageValueMeter()
previous_loss = 1e10
for epoch in range(opt.max_epoch):
loss_meter.reset()
for ii,data in tqdm(enumerate(train_dataloader)):
input = data[0].float().to(opt.device).squeeze(1)
target = input
input_ = input.view(input.size(0),input.size(1)*input.size(2),1)
weight = key_bernoulli_weights.repeat(input.size(0),1,1).to(opt.device)
data_i = t.bmm(weight,input_).view(input.size(0),int(opt.CR[0]*1024))
optimizer.zero_grad()
score = model(data_i)
loss = criterion(score,target)
loss.backward()
optimizer.step()
loss_meter.add(loss.item())
if(epoch%10==0):
print("loss value is:",loss_meter.value()[0])
val_psnr_ave = val(model,val_dataloader)
print("result on validation dataset is:",val_psnr_ave)
if loss_meter.value()[0] > previous_loss:
lr = lr * opt.lr_decay
for param_group in optimizer.param_groups:
param_group['lr'] = lr
previous_loss = loss_meter.value()[0]
model.save()
def val(model,dataloader):
model.eval()
utils_eval = utils.Evaluation()
for ii, data in tqdm(enumerate(dataloader)):
val_input = data[0].float().to(opt.device).squeeze(1)
val_input_ = val_input.view(val_input.size(0),val_input.size(1)*val_input.size(2),1)
weight = key_bernoulli_weights.repeat(val_input.size(0),1,1).to(opt.device)
data_i = t.bmm(weight,val_input_).view(val_input.size(0),int(opt.CR[0]*1024))
score = model(data_i).unsqueeze(1)
utils_eval.add(score,val_input.unsqueeze(1).float())
model.train()
psnr_ave = utils_eval.psnr_value()
return psnr_ave
def test(**kwargs):
opt._parse(kwargs)
model = models.Key_CSVideoNet(1, opt.Height, opt.Width, 1)
if opt.load_model_path:
model.load(opt.load_model_path)
print("model load success!")
#if t.cuda.device_count()>1:
# model = t.nn.DataParallel(model,device_ids=[0,1])
model.to(opt.device)
model.eval()
videos = [
os.path.join(opt.test_data_root, video)
for video in os.listdir(opt.test_data_root)
]
video_num = len(videos)
for item in videos:
ue = utils.Evaluation()
uv = utils.Video(item, opt.save_test_root, 1, opt.overlap, opt.Height)
frames = uv.frame_capture()
input = uv.frame_unfold(frames).float()
print(input.size())
frames_num = input.size(0)
block_num_h = input.size(1)
block_num_w = input.size(2)
output_ = t.zeros(frames_num, block_num_h, block_num_w, opt.Height,
opt.Width)
input_ = input.view(frames_num, block_num_h * block_num_w,
opt.Height * opt.Width).unsqueeze(3).to(opt.device)
for i in range(frames_num):
weight = key_bernoulli_weights.repeat(input_.size(1), 1,
1).to(opt.device)
#print(weight.size())
#print(input_.size())
input_t = input_[i] / 255.0
data_i = t.bmm(weight, input_t).view(input_.size(1), 1024)
output = model(data_i)
ue.add(
output.unsqueeze_(0),
input_t.view(input_.size(1), opt.Height,
opt.Width).unsqueeze_(0))
output_[i] = (output.view(block_num_h, block_num_w, opt.Height,
opt.Width)) * (255.0 / t.max(output))
frames_output = uv.frame_fold(output_)
uv.generate_video(frames_output)
print("the average PSNR is", ue.psnr_value())
'''
def val(model, dataloader):
model.eval()
utils_eval = utils.Evaluation()
for ii, data in tqdm(enumerate(dataloader)):
input = data
target = t.zeros(data[0].size(0), opt.seqLength, 32, 32).to(opt.device)
nonkey_m = t.zeros(input[0].size(0), opt.seqLength - 1,
int(1024 * opt.CR[1])).to(opt.device)
for i in range(opt.seqLength):
target[:, i, :, :] = data[i].float().to(opt.device).squeeze(1)
key_weight = key_bernoulli_weights.repeat(input[0].size(0), 1,
1).to(opt.device)
non_key_weight = non_key_bernoulli_weights.repeat(
input[1].size(0), 1, 1).to(opt.device)
key_m = t.bmm(
key_weight, input[0].float().to(opt.device).view(
input[0].size(0), 1024, 1)).view(input[0].size(0),
int(1024 * opt.CR[0]))
nonkey_m[:, 0, :] = t.bmm(
non_key_weight, input[1].view(input[1].size(0), 1024,
1).float().to(opt.device)).view(
input[1].size(0),
int(1024 * opt.CR[1]))
#nonkey_m[:,1,:] = t.bmm(non_key_weight,input[2].view(input[2].size(0),1024,1).float().to(opt.device)).view(input[2].size(0),int(1024*opt.CR[1]))
#nonkey_m[:,2,:] = t.bmm(non_key_weight,input[3].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,3,:] = t.bmm(non_key_weight,input[4].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,4,:] = t.bmm(non_key_weight,input[5].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,5,:] = t.bmm(non_key_weight,input[6].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,6,:] = t.bmm(non_key_weight,input[7].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,7,:] = t.bmm(non_key_weight,input[8].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,8,:] = t.bmm(non_key_weight,input[9].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
score = model(key_m, nonkey_m)
utils_eval.add(score, target)
model.train()
psnr_ave = utils_eval.psnr_value()
return psnr_ave
inputs=X_one_hot,
dtype=tf.float32,
sequence_length=Sequences)
outputs = tf.concat(values=outputs, axis=2)
model = tf.layers.dense(outputs, modelconfig.output_size, activation=None)
cost = tf.reduce_mean(
tf.nn.sparse_softmax_cross_entropy_with_logits(logits=model, labels=Y))
prediction = tf.argmax(model, axis=2)
optimizer = tf.train.AdamOptimizer(learning_rate=0.001).minimize(cost)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
eval = utils.Evaluation(type, modelconfig.epoch, 25)
print('\n------------ Training ------------ ')
last_time = time.time()
for epoch in range(modelconfig.epoch):
_, loss = sess.run([optimizer, cost],
feed_dict={
X: input_batch,
Y: target_batch,
Sequences: seq_lens,
keep_prob: 0.8
})
if epoch % 25 == 24:
result = sess.run(prediction,
feed_dict={
def run(self):
training_dataset, valid_dataset = data.make_sequences(
self.input, self.char2vec, self.output_char2vec, self.seq_length)
input_batch = training_dataset.input_batch
target_batch = training_dataset.target_batch
seq_lens = training_dataset.seq_lens
hidden_size = self.modelconfig.hidden_size
X = tf.placeholder(tf.int32, [None, self.seq_length]) # X data
X_onehot = tf.one_hot(
X,
self.modelconfig.input_size) # one hot: 1 -> 0 1 0 0 0 0 0 0 0 0
Y = tf.placeholder(tf.int32, [None, self.seq_length]) # Y label
Sequences = tf.placeholder(tf.int32, [None])
keep_prob = tf.placeholder(tf.float32)
if self.type == "multi":
print('\n****** MultiLayer LSTM Initialize ******')
with tf.variable_scope('cell_def'):
cell1 = tf.nn.rnn_cell.BasicLSTMCell(num_units=hidden_size,
state_is_tuple=True)
cell1 = tf.nn.rnn_cell.DropoutWrapper(
cell1, output_keep_prob=keep_prob)
cell2 = tf.nn.rnn_cell.BasicLSTMCell(num_units=hidden_size,
state_is_tuple=True)
cell2 = tf.nn.rnn_cell.DropoutWrapper(
cell2, output_keep_prob=keep_prob)
multi_cell = tf.nn.rnn_cell.MultiRNNCell([cell1, cell2])
with tf.variable_scope('rnn_def'):
outputs, _states = tf.nn.dynamic_rnn(multi_cell,
X_onehot,
dtype=tf.float32,
sequence_length=Sequences)
elif self.type == "bimul":
print('\n****** Bidirectional LSTM Initialize ******')
with tf.variable_scope('cell_def'):
forward = tf.nn.rnn_cell.BasicLSTMCell(num_units=hidden_size,
state_is_tuple=True)
forward = tf.nn.rnn_cell.DropoutWrapper(
forward, output_keep_prob=keep_prob)
backward = tf.nn.rnn_cell.BasicLSTMCell(num_units=hidden_size,
state_is_tuple=True)
backward = tf.nn.rnn_cell.DropoutWrapper(
backward, output_keep_prob=keep_prob)
with tf.variable_scope('rnn_def'):
outputs, states = tf.nn.bidirectional_dynamic_rnn(
forward,
backward,
inputs=X_onehot,
dtype=tf.float32,
sequence_length=Sequences)
outputs = tf.concat(values=outputs, axis=2)
model = tf.layers.dense(outputs,
self.modelconfig.output_size,
activation=None)
cost = tf.reduce_mean(
tf.nn.sparse_softmax_cross_entropy_with_logits(logits=model,
labels=Y))
prediction = tf.argmax(model, axis=2)
optimizer = tf.train.AdamOptimizer(learning_rate=0.005).minimize(cost)
# optimizer = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(cost)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
eval = utils.Evaluation(self.type, self.modelconfig.epoch, 25)
print('\n------------ Training ------------ ')
last_time = time.time()
for epoch in range(self.modelconfig.epoch):
_, loss = sess.run(
[optimizer, cost],
feed_dict={
X: input_batch,
Y: target_batch,
Sequences: seq_lens,
keep_prob: 0.8
})
if epoch % 25 == 24:
result = sess.run(prediction,
feed_dict={
X: valid_dataset.input_batch,
Y: valid_dataset.target_batch,
Sequences: valid_dataset.seq_lens,
keep_prob: 1
})
accuracy = tf.reduce_mean(
tf.cast(tf.equal(result, tf.cast(Y, tf.int64)),
tf.float32))
accuracy_ret = sess.run(
accuracy, feed_dict={Y: valid_dataset.target_batch})
speed = time.time() - last_time
print('Epoch:', '%04d ' % (epoch + 1), 'accuracy =',
'{:.6f} '.format(accuracy_ret), 'cost =',
'{:.6f}'.format(loss), 'speed =', '{:.2f}'.format(speed),
'sec')
last_time = time.time()
avg_p, avg_r, avg_f = utils.print_evaluation(
valid_dataset.target_batch, result,
self.output_char2vec.char_dict)
eval.set(epoch, accuracy_ret, loss, speed, avg_p, avg_r, avg_f)
print('')
print('\n------------ Testing ------------ ')
test_sentences = data.read_data("data/test/BHXX0035.txt", 30)
test_dataset, _ = data.make_sequences(test_sentences,
self.char2vec,
self.output_char2vec,
self.seq_length,
make_valid=False)
result = sess.run(prediction,
feed_dict={
X: test_dataset.input_batch,
Y: test_dataset.target_batch,
Sequences: test_dataset.seq_lens,
keep_prob: 1
})
accuracy = tf.reduce_mean(
tf.cast(tf.equal(result, tf.cast(Y, tf.int64)), tf.float32))
accuracy_ret = sess.run(accuracy,
feed_dict={Y: test_dataset.target_batch})
print('Accuracy =', '{:.6f}'.format(accuracy_ret))
for index, predict_sequence in enumerate(result):
target_output, prediction_output = data.compare_sentence(
self.output_char2vec, test_dataset.target_batch[index],
test_dataset.input_source[index], predict_sequence)
if index < 2:
print("target sentence: ", target_output[1])
print("prediction sentence:", prediction_output[1])
return eval
def train(**kwargs):
opt._parse(kwargs)
model = models.CSVideoNet(opt.CR, opt.seqLength, opt.Height, opt.Width, 1)
if opt.load_model_path:
model.load(opt.load_model_path, True)
'''
if t.cuda.device_count() > 1:
model = t.nn.DataParallel(model,device_ids=[0,1])
'''
model.to(opt.device)
train_data = UCF101(opt.train_data_root, train=True)
val_data = UCF101(opt.train_data_root, train=False)
train_dataloader = DataLoader(train_data,
opt.batch_size,
shuffle=True,
num_workers=opt.num_workers)
val_dataloader = DataLoader(val_data,
opt.batch_size,
shuffle=False,
num_workers=opt.num_workers)
criterion = t.nn.MSELoss()
lr = opt.lr
optimizer = t.optim.SGD(model.parameters(), lr=lr, momentum=opt.momentum)
utils_eval = utils.Evaluation()
loss_meter = meter.AverageValueMeter()
previous_loss = 1e10
for epoch in range(opt.max_epoch):
loss_meter.reset()
for ii, data in tqdm(enumerate(train_dataloader)):
#train model
input = data
target = t.zeros(data[0].size(0), opt.seqLength, 32,
32).to(opt.device)
nonkey_m = t.zeros(input[0].size(0), opt.seqLength - 1,
int(1024 * opt.CR[1])).to(opt.device)
for i in range(opt.seqLength):
target[:, i, :, :] = data[i].float().to(opt.device).squeeze(1)
key_weight = key_bernoulli_weights.repeat(input[0].size(0), 1,
1).to(opt.device)
non_key_weight = non_key_bernoulli_weights.repeat(
input[1].size(0), 1, 1).to(opt.device)
key_m = t.bmm(
key_weight, input[0].float().to(opt.device).view(
input[0].size(0), 1024, 1)).view(input[0].size(0),
int(1024 * opt.CR[0]))
nonkey_m[:, 0, :] = t.bmm(
non_key_weight, input[1].view(input[1].size(0), 1024,
1).float().to(opt.device)).view(
input[1].size(0),
int(1024 * opt.CR[1]))
#nonkey_m[:,1,:] = t.bmm(non_key_weight,input[2].view(input[2].size(0),1024,1).float().to(opt.device)).view(input[2].size(0),int(1024*opt.CR[1]))
#nonkey_m[:,2,:] = t.bmm(non_key_weight,input[3].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,3,:] = t.bmm(non_key_weight,input[4].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,4,:] = t.bmm(non_key_weight,input[5].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,5,:] = t.bmm(non_key_weight,input[6].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,6,:] = t.bmm(non_key_weight,input[7].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,7,:] = t.bmm(non_key_weight,input[8].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
#nonkey_m[:,8,:] = t.bmm(non_key_weight,input[9].view(opt.batch_size,1024,1).float().to(opt.device)).view(opt.batch_size,int(1024*opt.CR[1]))
optimizer.zero_grad()
score = model(key_m, nonkey_m)
loss = criterion(score, target)
loss.backward()
optimizer.step()
#t.nn.utils.clip_grad_norm(model.modules.parameters(),opt.gradient_clipping)
loss_meter.add(loss.item())
val_psnr_ave = val(model, val_dataloader)
print(epoch, "result on validation dataset is:", val_psnr_ave)
if loss_meter.value()[0] > previous_loss:
lr = lr * opt.lr_decay
for param_group in optimizer.param_groups:
param_group['lr'] = lr
previous_loss = loss_meter.value()[0]
model.save()
elif 'CIFAR' in args.dataset:
datasets = CIFAR.build_dataset_providers(args, strategy)
model = utils.load_model(args, datasets['num_classes'])
summary_writer = tf.summary.create_file_writer(args.train_path,
flush_millis=30000)
with summary_writer.as_default():
utils.save_code_and_augments(args)
total_train_time = 0
train_step, train_loss, train_accuracy, optimizer = op_utils.Optimizer(
args, model, strategy)
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.SUM)
Eval = utils.Evaluation(args, model, strategy, datasets['test'],
loss_object)
print('Training starts')
fine_tuning_time = 0
tic = time.time()
for step, data in enumerate(datasets['train']):
epoch = step // datasets['train_len']
lr = utils.scheduler(args, optimizer, epoch)
train_step(*data)
step += 1
if step % args.do_log == 0:
template = 'Global step {0:5d}: loss = {1:0.4f} ({2:1.3f} sec/step)'
train_time = time.time() - tic
print(
template.format(step,