def sample(load_dir: str, save_dir: str, use_gpu: bool) -> None:
'''
Sample the FantasyMapGAN for new maps.
Saves the generated images to `save_dir`.
Parameters
----------
load_dir: str
folder to load network weights from
save_dir: str
folder to save network weights to
use_gpu: bool
Set to true to run training on GPU, otherwise run on CPU
'''
# Network
model = Decoder()
model = model.eval()
if use_gpu:
model = model.cuda()
if load_dir:
fs = glob(os.path.join(load_dir, '*_dec.pth'))
fs.sort(key=os.path.getmtime)
model.load_state_dict(torch.load(fs[-1]))
# Generate
z = torch.randn((1, model.latent_dim))
x = model.forward(z)
# Save
save_path = os.path.join(save_dir, str(uuid.uuid1()) + '.png')
save_image(x.squeeze(), save_path)
class Tester(object):
def __init__(self, config):
self.content_images = glob.glob((config.exp_content_dir + '/*/*.jpg')) #+ '_resized/*'))
self.encoder = Encoder().cuda()
self.decoder = Decoder()
self.keyencoder = KeyEncoder().cuda()
self.decoder.load_state_dict(torch.load('./decoder.pth'))
self.decoder = self.decoder.cuda()
self.keyencoder.load_state_dict(torch.load('./key.pth'))
self.keyencoder = self.keyencoder.cuda()
if config.attention == 'soft':
self.AsyAtt = AsyAtt()
else:
self.AsyAtt = AsyAttHard()
S_path = os.path.join(config.style_dir, str(config.S))
style_images = glob.glob((S_path + '/*.jpg'))
s = Image.open(style_images[0])
s = trans(s).cuda()
self.style_image = s.unsqueeze(0)
self.style_target = torch.stack([s for i in range(config.batch_size)],0)
def test(self):
self.encoder.eval()
self.decoder.eval()
with torch.no_grad():
style_val = self.encoder(self.style_image)
style_key = self.keyencoder(style_val)
for filename in self.content_images:
name = str(filename).split("test_images")[-1][1:].replace("\\", "-")
name = name.replace("/", "-")
c = Image.open(filename)
c_tensor = trans(c).unsqueeze(0).cuda()
val = self.encoder(c_tensor)
key = self.keyencoder(val)
content_feature = self.AsyAtt(style_key[0], style_val[0], key, val)
out = self.decoder(content_feature)
out = denorm(out).to('cpu')[0]
c_tensor = denorm(c_tensor).to('cpu')[0]
if out.shape[1] > c_tensor.shape[1]:
c_tensor = torch.cat([c_tensor, torch.zeros([c_tensor.shape[0],out.shape[1]-c_tensor.shape[1],c_tensor.shape[2]])],1)
elif out.shape[1] < c_tensor.shape[1]:
out = torch.cat([out, torch.zeros([out.shape[0],c_tensor.shape[1]-out.shape[1],out.shape[2]])],1)
save_image(torch.cat([out, c_tensor], 2), os.path.join('./logs/test', name))
decoder = Decoder(emb_size, 2 * rnn_hid_size + z_hid_size, n_words, 1)
if reload_from > 0: # if using from previous data
embedder = torch.load(f='%slatentvar_%s/embedder.pckl' %
(model_dir, str(reload_from)))
encoder = torch.load(f='%slatentvar_%s/encoder.pckl' %
(model_dir, str(reload_from)))
hidvar = torch.load(f='%slatentvar_%s/hidvar.pckl' %
(model_dir, str(reload_from)))
decoder = torch.load(f='%slatentvar_%s/decoder.pckl' %
(model_dir, str(reload_from)))
if use_cuda:
embedder = embedder.cuda()
encoder = encoder.cuda()
hidvar = hidvar.cuda() #!!!!!
decoder = decoder.cuda()
TRAIN_FILE = input_dir + 'train.h5'
train_set = UbuntuDataset(TRAIN_FILE, max_seq_len=20)
train_data_loader = torch.utils.data.DataLoader(
dataset=train_set,
batch_size=batch_size,
shuffle=True,
num_workers=1 # multiple num_workers could introduce error (conflict?)
)
vocab = load_dict(input_dir + 'vocab.json')
train(embedder,
encoder,
hidvar,
decoder,
def train_dynamics(env, args, writer=None):
"""
Trains the Dynamics module. Supervised.
Arguments:
env: the initialized environment (rllab/gym)
args: input arguments
writer: initialized summary writer for tensorboard
"""
args.action_space = env.action_space
# Initialize models
enc = Encoder(env.observation_space.shape[0],
args.dim,
use_conv=args.use_conv)
dec = Decoder(env.observation_space.shape[0],
args.dim,
use_conv=args.use_conv)
d_module = D_Module(env.action_space.shape[0], args.dim, args.discrete)
if args.from_checkpoint is not None:
results_dict = torch.load(args.from_checkpoint)
enc.load_state_dict(results_dict['enc'])
dec.load_state_dict(results_dict['dec'])
d_module.load_state_dict(results_dict['d_module'])
all_params = chain(enc.parameters(), dec.parameters(),
d_module.parameters())
if args.transfer:
for p in enc.parameters():
p.requires_grad = False
for p in dec.parameters():
p.requires_grad = False
all_params = d_module.parameters()
optimizer = torch.optim.Adam(all_params,
lr=args.lr,
weight_decay=args.weight_decay)
if args.gpu:
enc = enc.cuda()
dec = dec.cuda()
d_module = d_module.cuda()
# Initialize datasets
val_loader = None
train_dataset = DynamicsDataset(args.train_set,
args.train_size,
batch=args.train_batch,
rollout=args.rollout)
val_dataset = DynamicsDataset(args.test_set,
5000,
batch=args.test_batch,
rollout=args.rollout)
val_loader = torch.utils.data.DataLoader(dataset=val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
results_dict = {
'dec_losses': [],
'forward_losses': [],
'inverse_losses': [],
'total_losses': [],
'enc': None,
'dec': None,
'd_module': None,
'd_init': None,
'args': args
}
total_action_taken = 0
correct_predicted_a_hat = 0
# create the mask here for re-weighting
dec_mask = None
if args.dec_mask is not None:
dec_mask = torch.ones(9)
game_vocab = dict([
(b, a)
for a, b in enumerate(sorted(env.game.all_possible_features()))
])
dec_mask[game_vocab['Agent']] = args.dec_mask
dec_mask[game_vocab['Goal']] = args.dec_mask
dec_mask = dec_mask.expand(args.batch_size, args.maze_length,
args.maze_length, 9).contiguous().view(-1)
dec_mask = Variable(dec_mask, requires_grad=False)
if args.gpu:
dec_mask = dec_mask.cuda()
for epoch in range(1, args.num_epochs + 1):
enc.train()
dec.train()
d_module.train()
if args.framework == "mazebase":
d_init.train()
# for measuring the accuracy
train_acc = 0
current_epoch_actions = 0
current_epoch_predicted_a_hat = 0
start = time.time()
for i, (states, target_actions) in enumerate(train_loader):
optimizer.zero_grad()
if args.framework != "mazebase":
forward_loss, inv_loss, dec_loss, recon_loss, model_loss, _, _ = forward_planning(
i, states, target_actions, enc, dec, d_module, args)
else:
forward_loss, inv_loss, dec_loss, recon_loss, model_loss, current_epoch_predicted_a_hat, current_epoch_actions = multiple_forward(
i, states, target_actions, enc, dec, d_module, args,
d_init, dec_mask)
loss = forward_loss + args.inv_loss_coef * inv_loss + \
args.dec_loss_coef * dec_loss
if i % args.log_interval == 0:
log(
'Epoch [{}/{}]\tIter [{}/{}]\t'.format(
epoch, args.num_epochs, i+1, len(
train_dataset)//args.batch_size) + \
'Time: {:.2f}\t'.format(time.time() - start) + \
'Decoder Loss: {:.2f}\t'.format(dec_loss.data[0]) + \
'Forward Loss: {:.2f}\t'.format(forward_loss.data[0] ) + \
'Inverse Loss: {:.2f}\t'.format(inv_loss.data[0]) + \
'Loss: {:.2f}\t'.format(loss.data[0]))
results_dict['dec_losses'].append(dec_loss.data[0])
results_dict['forward_losses'].append(forward_loss.data[0])
results_dict['inverse_losses'].append(inv_loss.data[0])
results_dict['total_losses'].append(loss.data[0])
# write the summaries here
if writer:
writer.add_scalar('dynamics/total_loss', loss.data[0],
epoch)
writer.add_scalar('dynamics/decoder', dec_loss.data[0],
epoch)
writer.add_scalar('dynamics/reconstruction_loss',
recon_loss.data[0], epoch)
writer.add_scalar('dynamics/next_state_prediction_loss',
model_loss.data[0], epoch)
writer.add_scalar('dynamics/inv_loss', inv_loss.data[0],
epoch)
writer.add_scalar('dynamics/forward_loss',
forward_loss.data[0], epoch)
writer.add_scalars(
'dynamics/all_losses', {
"total_loss": loss.data[0],
"reconstruction_loss": recon_loss.data[0],
"next_state_prediction_loss": model_loss.data[0],
"decoder_loss": dec_loss.data[0],
"inv_loss": inv_loss.data[0],
"forward_loss": forward_loss.data[0],
}, epoch)
loss.backward()
correct_predicted_a_hat += current_epoch_predicted_a_hat
total_action_taken += current_epoch_actions
# does it not work at all without grad clipping ?
torch.nn.utils.clip_grad_norm(all_params, args.max_grad_norm)
optimizer.step()
# maybe add the generated image to add the logs
# writer.add_image()
# Run validation
if val_loader is not None:
enc.eval()
dec.eval()
d_module.eval()
forward_loss, inv_loss, dec_loss = 0, 0, 0
for i, (states, target_actions) in enumerate(val_loader):
f_loss, i_loss, d_loss, _, _, _, _ = forward_planning(
i, states, target_actions, enc, dec, d_module, args)
forward_loss += f_loss
inv_loss += i_loss
dec_loss += d_loss
loss = forward_loss + args.inv_loss_coef * inv_loss + \
args.dec_loss_coef * dec_loss
if writer:
writer.add_scalar('val/forward_loss', forward_loss.data[0] / i,
epoch)
writer.add_scalar('val/inverse_loss', inv_loss.data[0] / i,
epoch)
writer.add_scalar('val/decoder_loss', dec_loss.data[0] / i,
epoch)
log(
'[Validation]\t' + \
'Decoder Loss: {:.2f}\t'.format(dec_loss.data[0] / i) + \
'Forward Loss: {:.2f}\t'.format(forward_loss.data[0] / i) + \
'Inverse Loss: {:.2f}\t'.format(inv_loss.data[0] / i) + \
'Loss: {:.2f}\t'.format(loss.data[0] / i))
if epoch % args.checkpoint == 0:
results_dict['enc'] = enc.state_dict()
results_dict['dec'] = dec.state_dict()
results_dict['d_module'] = d_module.state_dict()
if args.framework == "mazebase":
results_dict['d_init'] = d_init.state_dict()
torch.save(
results_dict,
os.path.join(args.out, 'dynamics_module_epoch%s.pt' % epoch))
log('Saved model %s' % epoch)
results_dict['enc'] = enc.state_dict()
results_dict['dec'] = dec.state_dict()
results_dict['d_module'] = d_module.state_dict()
torch.save(results_dict,
os.path.join(args.out, 'dynamics_module_epoch%s.pt' % epoch))
print(os.path.join(args.out, 'dynamics_module_epoch%s.pt' % epoch))
def main(_):
# Load the configuration file.
with open(FLAGS.config, 'r') as f:
config = yaml.load(f)
# Create the checkpoint directory if it does not already exist.
ckpt_dir = os.path.join(config['data']['ckpt'], config['experiment_name'])
if not os.path.exists(ckpt_dir):
os.mkdir(ckpt_dir)
# Check if a pre-existing configuration file exists and matches the current
# configuration. Otherwise save a copy of the configuration to the
# checkpoint directory.
prev_config_path = os.path.join(ckpt_dir, 'config.yaml')
if os.path.exists(prev_config_path):
with open(prev_config_path, 'r') as f:
prev_config = yaml.load(f)
assert config == prev_config
else:
shutil.copyfile(FLAGS.config, prev_config_path)
# Load the vocabularies.
src_vocab = Vocab.load(config['data']['src']['vocab'])
tgt_vocab = Vocab.load(config['data']['tgt']['vocab'])
# Load the training and dev datasets.
train_data = ShakespeareDataset('train', config, src_vocab, tgt_vocab)
dev_data = ShakespeareDataset('dev', config, src_vocab, tgt_vocab)
# Build the model.
src_vocab_size = len(src_vocab)
tgt_vocab_size = len(tgt_vocab)
encoder = Encoder(src_vocab_size, config['model']['embedding_dim'])
decoder = Decoder(tgt_vocab_size, config['model']['embedding_dim'])
if torch.cuda.is_available():
encoder = encoder.cuda()
decoder = decoder.cuda()
# Define the loss function + optimizer.
loss_weights = torch.ones(decoder.tgt_vocab_size)
loss_weights[0] = 0
if torch.cuda.is_available():
loss_weights = loss_weights.cuda()
criterion = torch.nn.NLLLoss(loss_weights)
learning_rate = config['training']['learning_rate']
encoder_optimizer = torch.optim.Adam(encoder.parameters(),
lr=learning_rate)
decoder_optimizer = torch.optim.Adam(decoder.parameters(),
lr=learning_rate)
# Restore saved model (if one exists).
ckpt_path = os.path.join(ckpt_dir, 'model.pt')
if os.path.exists(ckpt_path):
print('Loading checkpoint: %s' % ckpt_path)
ckpt = torch.load(ckpt_path)
epoch = ckpt['epoch']
encoder.load_state_dict(ckpt['encoder'])
decoder.load_state_dict(ckpt['decoder'])
encoder_optimizer.load_state_dict(ckpt['encoder_optimizer'])
decoder_optimizer.load_state_dict(ckpt['decoder_optimizer'])
else:
epoch = 0
train_log_string = '%s :: Epoch %i :: Iter %i / %i :: train loss: %0.4f'
dev_log_string = '\n%s :: Epoch %i :: dev loss: %0.4f'
while epoch < config['training']['num_epochs']:
# Main training loop.
train_loss = []
sampler = RandomSampler(train_data)
for i, train_idx in enumerate(sampler):
src, tgt = train_data[train_idx]
# Clear gradients
encoder_optimizer.zero_grad()
decoder_optimizer.zero_grad()
# Feed inputs one by one from src into encoder (in reverse).
src_length = src.size()[0]
hidden = None
for j in reversed(range(src_length)):
encoder_output, hidden = encoder(src[j], hidden)
# Feed desired outputs one by one from tgt into decoder
# and measure loss.
tgt_length = tgt.size()[0]
loss = 0
for j in range(tgt_length - 1):
decoder_output, hidden = decoder(tgt[j], hidden)
loss += criterion(decoder_output, tgt[j + 1])
# Backpropagate the loss and update the model parameters.
loss.backward()
encoder_optimizer.step()
decoder_optimizer.step()
train_loss.append(loss.data.cpu())
# Every once and a while check on the loss
if ((i + 1) % 100) == 0:
print(train_log_string %
(datetime.now(), epoch, i + 1, len(train_data),
np.mean(train_loss)),
end='\r')
train_loss = []
# Evaluation loop.
dev_loss = []
for src, tgt in dev_data:
# Feed inputs one by one from src into encoder.
src_length = src.size()[0]
hidden = None
for j in reversed(range(src_length)):
encoder_output, hidden = encoder(src[j], hidden)
# Feed desired outputs one by one from tgt into decoder
# and measure loss.
tgt_length = tgt.size()[0]
loss = 0
for j in range(tgt_length - 1):
decoder_output, hidden = decoder(tgt[j], hidden)
loss += criterion(decoder_output, tgt[j + 1])
dev_loss.append(loss.data.cpu())
print(dev_log_string % (datetime.now(), epoch, np.mean(dev_loss)))
state_dict = {
'epoch': epoch,
'encoder': encoder.state_dict(),
'decoder': decoder.state_dict(),
'encoder_optimizer': encoder_optimizer.state_dict(),
'decoder_optimizer': decoder_optimizer.state_dict()
}
torch.save(state_dict, ckpt_path)
epoch += 1
def main(_):
# Load the configuration file.
with open(FLAGS.config, 'r') as f:
config = yaml.load(f)
# Load the vocabularies.
src_vocab = Vocab.load(config['data']['src']['vocab'])
tgt_vocab = Vocab.load(config['data']['tgt']['vocab'])
# Load the training and dev datasets.
test_data = ShakespeareDataset('test', config, src_vocab, tgt_vocab)
# Restore the model.
src_vocab_size = len(src_vocab)
tgt_vocab_size = len(tgt_vocab)
encoder = Encoder(src_vocab_size, config['model']['embedding_dim'],
config['model']['bidirection'],
config['model']['dropout'], config['model']['layer'],
config['model']['mode'])
decoder = Decoder(tgt_vocab_size, config['model']['embedding_dim'],
config['model']['bidirection'],
config['model']['dropout'], config['model']['layer'],
config['model']['mode'])
if torch.cuda.is_available():
encoder = encoder.cuda()
decoder = decoder.cuda()
ckpt_path = os.path.join(config['data']['ckpt'], config['experiment_name'],
'model.pt')
if os.path.exists(ckpt_path):
print('Loading checkpoint: %s' % ckpt_path)
ckpt = torch.load(ckpt_path)
encoder.load_state_dict(ckpt['encoder'])
decoder.load_state_dict(ckpt['decoder'])
else:
print('Unable to find checkpoint. Terminating.')
sys.exit(1)
encoder.eval()
decoder.eval()
# Initialize translator.
greedy_translator = GreedyTranslator(encoder, decoder, tgt_vocab)
# Qualitative evaluation - print translations for first couple sentences in
# test corpus.
for i in range(10):
src, tgt = test_data[i]
translation = greedy_translator(src)
src_sentence = [src_vocab.id2word(id) for id in src.data.cpu().numpy()]
tgt_sentence = [tgt_vocab.id2word(id) for id in tgt.data.cpu().numpy()]
translated_sentence = [tgt_vocab.id2word(id) for id in translation]
print('---')
print('Source: %s' % ' '.join(src_sentence))
print('Ground truth: %s' % ' '.join(tgt_sentence))
print('Model output: %s' % ' '.join(translated_sentence))
print('---')
# Quantitative evaluation - compute corpus level BLEU scores.
hypotheses = []
references = []
for src, tgt in test_data:
translation = greedy_translator(src)
tgt_sentence = [tgt_vocab.id2word(id) for id in tgt.data.cpu().numpy()]
translated_sentence = [tgt_vocab.id2word(id) for id in translation]
# Remove start and end of sentence tokens.
tgt_sentence = tgt_sentence[1:-1]
translated_sentence = translated_sentence[1:-1]
hypotheses.append(tgt_sentence)
references.append([translated_sentence])
print("Corpus BLEU score: %0.4f" % corpus_bleu(references, hypotheses))
class Trainer:
def __init__(self, driving, target, time_step, split, lr):
self.dataset = DataSet(driving, target, time_step, split)
f = open('dataset_obj.txt', 'wb')
pickle.dump(self.dataset, f)
f.close()
print('save model finish!!!!!!!!!!!!!!!!!!')
# f = open('dataset_obj.txt','rb')
# self.dataset = pickle.load(f)
# f.close()
self.encoder = Encoder(input_size=self.dataset.get_num_features(),
hidden_size=ENCODER_HIDDEN_SIZE,
T=time_step)
self.decoder = Decoder(encoder_hidden_size=ENCODER_HIDDEN_SIZE,
decoder_hidden_size=DECODER_HIDDEN_SIZE,
T=time_step)
if torch.cuda.is_available():
# print('tocuda')
self.encoder = self.encoder.cuda()
self.decoder = self.decoder.cuda()
self.encoder_optim = optim.Adam(self.encoder.parameters(), lr)
self.decoder_optim = optim.Adam(self.decoder.parameters(), lr)
self.loss_func = nn.CrossEntropyLoss()
self.train_size, self.validation_size, self.test_size = self.dataset.get_size(
)
self.best_dev_acc = 0.0
def get_accuracy(self, truth, pred):
assert len(truth) == len(pred)
right = (truth == pred).sum()
return right / len(truth)
def train_minibatch(self, num_epochs, batch_size, interval):
train_acc_list = []
dev_acc_list = []
train_loss_list = []
dev_loss_list = []
x_train, y_train, y_seq_train = self.dataset.get_train_set()
# print(x_train.shape)
for epoch in range(num_epochs):
print('Start epoch {}'.format(epoch))
i = 0
loss_sum = 0
pred_res_total = []
while i < self.train_size:
self.encoder_optim.zero_grad()
self.decoder_optim.zero_grad()
batch_end = i + batch_size
if (batch_end >= self.train_size):
batch_end = self.train_size
var_x = self.to_variable(x_train[i:batch_end])
# var_y = self.to_variable(y_train[i: batch_end])
var_y = Variable(
torch.from_numpy(y_train[i:batch_end]).long()).cuda()
var_y_seq = self.to_variable(y_seq_train[i:batch_end])
if var_x.dim() == 2:
var_x = var_x.unsqueeze(2)
code = self.encoder(var_x)
y_res = self.decoder(code, var_y_seq)
loss = self.loss_func(y_res, var_y)
if i == 0:
print("y_res:", y_res)
print("var_y:", var_y)
loss.backward()
self.encoder_optim.step()
self.decoder_optim.step()
loss_sum += loss.item()
# update the i
i = batch_end
pred_y = y_res.data.cpu()
# print('see what the pred and truth')
# print('y_res:',y_res.shape,' : ',y_res)
# print('var_y:',var_y.shape,' : ',var_y)
pred_y = torch.max(F.softmax(pred_y, dim=1), 1)[1]
#
# print('pred_y:',pred_y)
# print('var_y',var_y)
pred_res_total.extend(pred_y)
# if i%50 == 0:
# print(' finish {0:.2f}/100'.format(i/self.train_size))
acc = self.get_accuracy(y_train, np.array(pred_res_total))
print(
'epoch [%d] finished, the average loss is %.2f, accuracy is %.1f'
% (epoch, loss_sum, acc * 100))
dev_acc, dev_loss = self.test(batch_size)
print('dev_acc is %.2f' % (dev_acc * 100))
train_acc_list.append(acc)
dev_acc_list.append(dev_acc)
train_loss_list.append(loss_sum)
dev_loss_list.append(dev_loss)
if dev_acc > self.best_dev_acc:
torch.save(
self.encoder.state_dict(),
'D:\Projects\\stock_predict\\models\\encoder_best.model')
torch.save(
self.decoder.state_dict(),
'D:\Projects\\stock_predict\\models\\decoder_best.model')
self.best_dev_acc = dev_acc
test_acc, test_loss = self.test(batch_size, True)
print('test_accuracy: %.1f' % (test_acc * 100))
return train_acc_list, dev_acc_list, train_loss_list, dev_loss_list
def test(self, batch_size, is_test=False):
if not is_test:
x, y, y_seq = self.dataset.get_validation_set()
else:
x, y, y_seq = self.dataset.get_test_set()
i = 0
res = []
length = len(y)
loss_sum = 0
while i < length:
batch_end = i + batch_size
if batch_end >= length:
batch_end = length
var_x = self.to_variable(x[i:batch_end])
var_y = Variable(torch.from_numpy(y[i:batch_end]).long()).cuda()
# var_y = self.to_variable(y_test[i: batch_end])
var_y_seq = self.to_variable(y_seq[i:batch_end])
if var_x.dim() == 2:
var_x = var_x.unsqueeze(2)
# to encoder get encoder output
code = self.encoder(var_x)
# to decoder get classification
y_res = self.decoder(code, var_y_seq)
loss = self.loss_func(y_res, var_y)
loss_sum += loss.item()
pred_y = y_res.data.cpu()
pred_y = torch.max(pred_y, 1)[1]
res.extend(pred_y)
i = batch_end
res = np.array(res)
return self.get_accuracy(y, res), loss_sum
def load_model(self, encoder_path, decoder_path):
self.encoder.load_state_dict(
torch.load(encoder_path,
map_location=lambda storage, loc: storage))
self.decoder.load_state_dict(
torch.load(decoder_path,
map_location=lambda storage, loc: storage))
def to_variable(self, x):
if torch.cuda.is_available():
# print("var to cuda")
return Variable(torch.from_numpy(x).float()).cuda()
else:
return Variable(torch.from_numpy(x).float())
def draw_plot(self, train_list, dev_list, acc=True):
plt.plot(np.array(train_list))
plt.plot(np.array(dev_list))
if acc:
plt.title('model acc')
plt.ylabel('accuracy')
else:
plt.title('model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'validation'], loc='upper left')
plt.show()
def main():
# Praise argparser!
parser = argparse.ArgumentParser(
description=
"Inference script for performing joint tasks on ATIS datasets.")
parser.add_argument("--train_path",
type=str,
help="path of train dataset.")
parser.add_argument("--test_path", type=str, help="path of test dataset.")
parser.add_argument("--model_dir",
type=str,
default="./models/",
help='path for saved trained models.')
parser.add_argument('--max_length',
type=int,
default=60,
help='max sequence length')
parser.add_argument('--embedding_size',
type=int,
default=100,
help='dimension of word embedding vectors')
parser.add_argument('--hidden_size',
type=int,
default=50,
help='dimension of lstm hidden states')
args = parser.parse_args()
# Load data
print("Loading data...")
_, word2index, tag2index, intent2index = preprocessing(
args.train_path, args.max_length)
index2tag = {v: k for k, v in tag2index.items()}
index2intent = {v: k for k, v in intent2index.items()}
# Load model
print("Loading model...")
encoder = Encoder(len(word2index), args.embedding_size, args.hidden_size)
decoder = Decoder(len(tag2index), len(intent2index),
len(tag2index) // 3, args.hidden_size * 2)
encoder.load_state_dict(
torch.load(os.path.join(args.model_dir, 'jointnlu-encoder.pkl'),
map_location=None if USE_CUDA else "cpu"))
decoder.load_state_dict(
torch.load(os.path.join(args.model_dir, 'jointnlu-decoder.pkl'),
map_location=None if USE_CUDA else "cpu"))
if USE_CUDA:
encoder = encoder.cuda()
decoder = decoder.cuda()
# Switch to evaluation mode
encoder.eval()
decoder.eval()
# Preprocess test data
test = open(args.test_path, "r").readlines()
test = [t[:-1] for t in test]
test = [[
t.split("\t")[0].split(" "),
t.split("\t")[1].split(" ")[:-1],
t.split("\t")[1].split(" ")[-1]
] for t in test]
test = [
[t[0][1:-1], t[1][1:], t[2].split("#")[0]] for t in test
] # Note here I split embedded multiple labels into separate labels and get the first one.
# This could lower error rate.
slot_f1 = []
intent_err = []
# Test cases.
for index in range(len(test)):
test_raw = test[index][0]
test_in = prepare_sequence(test_raw, word2index).to("cpu")
test_mask = Variable(
torch.BoolTensor(tuple(map(
lambda s: s == 0,
test_in.data)))).cuda() if USE_CUDA else Variable(
torch.BoolTensor(tuple(map(lambda s: s == 0,
test_in.data)))).view(1, -1)
if USE_CUDA:
start_decode = Variable(
torch.LongTensor([[word2index['<SOS>']] * 1
])).cuda().transpose(1, 0)
else:
start_decode = Variable(
torch.LongTensor([[word2index['<SOS>']] * 1])).transpose(1, 0)
output, hidden_c = encoder(test_in.unsqueeze(0),
test_mask.unsqueeze(0))
tag_score, intent_score = decoder(start_decode, hidden_c, output,
test_mask)
v, i = torch.max(tag_score, 1)
slot_pred = list(map(lambda ii: index2tag[ii], i.data.tolist()))
slot_gt = test[index][1]
# Calculate f1_micro with sklearn. Pretty handy.
slot_f1.append(f1_score(slot_gt, slot_pred, average="micro"))
v, i = torch.max(intent_score, 1)
intent_pred = index2intent[i.data.tolist()[0]]
intent_gt = test[index][2]
if intent_pred != intent_gt:
intent_err.append([test[index][0], intent_gt, intent_pred])
# Print our results.
print("Input Sentence\t: ", *test[index][0])
print("Truth\t\t: ", *slot_gt)
print("Prediction\t: ", *slot_pred)
print("Truth\t\t: ", intent_gt)
print("Prediction\t: ", intent_pred)
print()
# Print out everything I need to finish my report.
# print("Got slot err ", len(slot_err[0]))
# print(*slot_err, sep="\n")
print("Got intent err ", len(intent_err))
print("--- BEGIN ERR PRINT ---")
for case in intent_err:
print("Input : ", *case[0])
print("Truth : ", case[1])
print("Predict: ", case[2])
print()
print("--- ENDOF ERR PRINT ---")
print("Total ", len(test))
print("Slot f1_micro avg %f" % np.average(slot_f1))
print("Intent acc %f" % (1 - len(intent_err) / len(test)))
batch_size=NUM_BATCH,
shuffle=True)
def cuda_tensors(obj):
for attr in dir(obj):
value = getattr(obj, attr)
if isinstance(value, torch.Tensor):
setattr(obj, attr, value.cuda())
enc = Encoder()
dec = Decoder()
if CUDA:
enc.cuda()
dec.cuda()
cuda_tensors(enc)
cuda_tensors(dec)
optimizer = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()),
lr=LEARNING_RATE,
betas=(BETA1, 0.999))
def elbo(q, p, alpha=0.1):
if NUM_SAMPLES is None:
return probtorch.objectives.montecarlo.elbo(q,
p,
sample_dim=None,
batch_dim=0,
alpha=alpha)
class Image_Captioning:
def __init__(self):
parser = argparse.ArgumentParser(description='Image Captioning')
parser.add_argument('--root',
default='../../../cocodataset/',
type=str)
parser.add_argument('--crop_size', default=224, type=int)
parser.add_argument('--epochs', default=100, type=int)
parser.add_argument('--lr', default=1e-4, type=float)
parser.add_argument('--batch_size', default=128, help='')
parser.add_argument('--num_workers', default=4, type=int)
parser.add_argument('--embed_dim', default=256, type=int)
parser.add_argument('--hidden_size', default=512, type=int)
parser.add_argument('--num_layers', default=1, type=int)
parser.add_argument('--model_path', default='./model/', type=str)
parser.add_argument('--vocab_path', default='./vocab/', type=str)
parser.add_argument('--save_step', default=1000, type=int)
self.args = parser.parse_args()
self.Multi_GPU = False
# if torch.cuda.device_count() > 1:
# print('Multi GPU Activate!')
# print('Using GPU :', int(torch.cuda.device_count()))
# self.Multi_GPU = True
os.makedirs(self.args.model_path, exist_ok=True)
transform = transforms.Compose([
transforms.RandomCrop(self.args.crop_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
with open(self.args.vocab_path + 'vocab.pickle', 'rb') as f:
data = pickle.load(f)
self.vocab = data
self.DataLoader = get_dataloader(root=self.args.root,
transform=transform,
shuffle=True,
batch_size=self.args.batch_size,
num_workers=self.args.num_workers,
vocab=self.vocab)
self.Encoder = Encoder(embed_dim=self.args.embed_dim)
self.Decoder = Decoder(embed_dim=self.args.embed_dim,
hidden_size=self.args.hidden_size,
vocab_size=len(self.vocab),
num_layers=self.args.num_layers)
# print(self.Encoder)
# print(self.Decoder)
def train(self):
if self.Multi_GPU:
self.Encoder = torch.nn.DataParallel(self.Encoder)
self.Decoder = torch.nn.DataParallel(self.Decoder)
parameters = list(self.Encoder.module.fc.parameters()) + list(
self.Encoder.module.BN.parameters()) + list(
self.Decoder.parameters())
else:
parameters = list(self.Encoder.fc.parameters()) + list(
self.Encoder.BN.parameters()) + list(self.Decoder.parameters())
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(parameters, lr=self.args.lr)
self.Encoder.cuda()
self.Decoder.cuda()
self.Encoder.train()
self.Decoder.train()
print('-' * 100)
print('Now Training')
print('-' * 100)
for epoch in range(self.args.epochs):
total_loss = 0
for batch_idx, (image, captions,
lengths) in enumerate(self.DataLoader):
optimizer.zero_grad()
image, captions = image.cuda(), captions.cuda()
targets = pack_padded_sequence(captions,
lengths,
batch_first=True)[0]
if self.Multi_GPU:
img_features = nn.parallel.DataParallel(
self.Encoder, image)
outputs = nn.parallel.DataParallel(
self.Decoder, (img_features, captions, lengths))
else:
img_features = self.Encoder(image)
outputs = self.Decoder(img_features, captions, lengths)
loss = criterion(outputs, targets)
total_loss += loss.item()
loss.backward()
optimizer.step()
if batch_idx % 30 == 0:
print('Epoch : {}, Step : [{}/{}], Step Loss : {:.4f}'.
format(epoch, batch_idx, len(self.DataLoader),
loss.item()))
print('Epoch : [{}/{}], Total loss : {:.4f}'.format(
epoch, self.args.epochs, total_loss / len(self.DataLoader)))
print('Now saving the models')
torch.save(
self.Encoder.state_dict(),
self.args.model_path + 'Encoder-{}.ckpt'.format(self.args.epochs))
torch.save(
self.Decoder.state_dict(),
self.args.model_path + 'Decoder-{}.ckpt'.format(self.args.epochs))
# denoise_enc = torch.load(args.edge_dir)
denoise_enc.load_state_dict(torch.load(args.denoise_dir))
edge_enc.load_state_dict(torch.load(args.edge_dir))
log('load trained model')
denoise_enc.eval()
edge_enc.eval()
model.train()
# criterion = nn.MSELoss(reduction = 'sum') # PyTorch 0.4.1
# criterion = sum_squared_error()
criterion = nn.MSELoss()
if cuda:
model = model.cuda()
denoise_enc = denoise_enc.cuda()
edge_enc = edge_enc.cuda()
# device_ids = [0]
# model = nn.DataParallel(model, device_ids=device_ids).cuda()
criterion = criterion.cuda()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = MultiStepLR(optimizer, milestones=[30, 60, 90],
gamma=0.2) # learning rates
for epoch in range(initial_epoch, n_epoch):
scheduler.step(epoch) # step to the learning rate in this epcoh
xs = dg.datagenerator(data_dir=args.train_data)
xs = xs.astype('float32') / 255.0
xs = torch.from_numpy(xs.transpose(
class Solver(object):
def __init__(self, hps, data_loader, log_dir='./log/'):
self.hps = hps
self.data_loader = data_loader
self.model_kept = []
self.max_keep = 20
self.build_model()
self.logger = Logger(log_dir)
def build_model(self):
hps = self.hps
ns = self.hps.ns
emb_size = self.hps.emb_size
self.Encoder = Encoder(ns=ns, dp=hps.enc_dp)
self.Decoder = Decoder(ns=ns, c_a=hps.n_speakers, emb_size=emb_size)
self.Generator = Decoder(ns=ns, c_a=hps.n_speakers, emb_size=emb_size)
self.LatentDiscriminator = LatentDiscriminator(ns=ns, dp=hps.dis_dp)
self.PatchDiscriminator = PatchDiscriminator(ns=ns,
n_class=hps.n_speakers)
if torch.cuda.is_available():
self.Encoder.cuda()
self.Decoder.cuda()
self.Generator.cuda()
self.LatentDiscriminator.cuda()
self.PatchDiscriminator.cuda()
betas = (0.5, 0.9)
params = list(self.Encoder.parameters()) + list(
self.Decoder.parameters())
self.ae_opt = optim.Adam(params, lr=self.hps.lr, betas=betas)
self.gen_opt = optim.Adam(self.Generator.parameters(),
lr=self.hps.lr,
betas=betas)
self.lat_opt = optim.Adam(self.LatentDiscriminator.parameters(),
lr=self.hps.lr,
betas=betas)
self.patch_opt = optim.Adam(self.PatchDiscriminator.parameters(),
lr=self.hps.lr,
betas=betas)
def save_model(self, model_path, iteration, enc_only=True):
if not enc_only:
all_model = {
'encoder': self.Encoder.state_dict(),
'decoder': self.Decoder.state_dict(),
'generator': self.Generator.state_dict(),
'latent_discriminator': self.LatentDiscriminator.state_dict(),
'patch_discriminator': self.PatchDiscriminator.state_dict(),
}
else:
all_model = {
'encoder': self.Encoder.state_dict(),
'decoder': self.Decoder.state_dict(),
'generator': self.Generator.state_dict(),
}
new_model_path = '{}-{}'.format(model_path, iteration)
with open(new_model_path, 'wb') as f_out:
torch.save(all_model, f_out)
self.model_kept.append(new_model_path)
if len(self.model_kept) >= self.max_keep:
os.remove(self.model_kept[0])
self.model_kept.pop(0)
def load_model(self, model_path, enc_only=True):
print('load model from {}'.format(model_path))
with open(model_path, 'rb') as f_in:
all_model = torch.load(f_in)
self.Encoder.load_state_dict(all_model['encoder'])
self.Decoder.load_state_dict(all_model['decoder'])
#self.Genrator.load_state_dict(all_model['generator'])
if not enc_only:
self.LatentDiscriminator.load_state_dict(
all_model['latent_discriminator'])
self.PatchDiscriminator.load_state_dict(
all_model['patch_discriminator'])
def set_eval(self):
self.Encoder.eval()
self.Decoder.eval()
self.Generator.eval()
#self.LatentDiscriminator.eval()
def test_step(self, x, c):
self.set_eval()
x = to_var(x).permute(0, 2, 1)
enc = self.Encoder(x)
x_tilde = self.Decoder(enc, c)
return x_tilde.data.cpu().numpy()
def permute_data(self, data):
C = [to_var(c, requires_grad=False) for c in data[:2]]
X = [to_var(x).permute(0, 2, 1) for x in data[2:]]
return C, X
def sample_c(self, size):
c_sample = Variable(torch.multinomial(torch.ones(8),
num_samples=size,
replacement=True),
requires_grad=False)
c_sample = c_sample.cuda() if torch.cuda.is_available() else c_sample
return c_sample
def cal_acc(self, logits, y_true):
_, ind = torch.max(logits, dim=1)
acc = torch.sum(
(ind == y_true).type(torch.FloatTensor)) / y_true.size(0)
return acc
def encode_step(self, *args):
enc_list = []
for x in args:
enc = self.Encoder(x)
enc_list.append(enc)
return tuple(enc_list)
def decode_step(self, enc, c):
x_tilde = self.Decoder(enc, c)
return x_tilde
def latent_discriminate_step(self,
enc_i_t,
enc_i_tk,
enc_i_prime,
enc_j,
is_dis=True):
same_pair = torch.cat([enc_i_t, enc_i_tk], dim=1)
diff_pair = torch.cat([enc_i_prime, enc_j], dim=1)
if is_dis:
same_val = self.LatentDiscriminator(same_pair)
diff_val = self.LatentDiscriminator(diff_pair)
w_dis = torch.mean(same_val - diff_val)
gp = calculate_gradients_penalty(self.LatentDiscriminator,
same_pair, diff_pair)
return w_dis, gp
else:
diff_val = self.LatentDiscriminator(diff_pair)
loss_adv = -torch.mean(diff_val)
return loss_adv
def patch_discriminate_step(self, x, x_tilde, cal_gp=True):
# w-distance
D_real, real_logits = self.PatchDiscriminator(x, classify=True)
D_fake, fake_logits = self.PatchDiscriminator(x_tilde, classify=True)
w_dis = torch.mean(D_real - D_fake)
if cal_gp:
gp = calculate_gradients_penalty(self.PatchDiscriminator, x,
x_tilde)
return w_dis, real_logits, fake_logits, gp
else:
return w_dis, real_logits, fake_logits
# backup
#def classify():
# # aux clssify loss
# criterion = nn.NLLLoss()
# c_loss = criterion(real_logits, c) + criterion(fake_logits, c_sample)
# real_acc = self.cal_acc(real_logits, c)
# fake_acc = self.cal_acc(fake_logits, c_sample)
def train(self, model_path, flag='train'):
# load hyperparams
hps = self.hps
for iteration in range(hps.iters):
# calculate current alpha
if iteration + 1 < hps.lat_sched_iters and iteration >= hps.enc_pretrain_iters:
current_alpha = hps.alpha_enc * (
iteration + 1 - hps.enc_pretrain_iters) / (
hps.lat_sched_iters - hps.enc_pretrain_iters)
else:
current_alpha = 0
if iteration >= hps.enc_pretrain_iters:
n_latent_steps = hps.n_latent_steps \
if iteration > hps.enc_pretrain_iters else hps.dis_pretrain_iters
for step in range(n_latent_steps):
#===================== Train latent discriminator =====================#
data = next(self.data_loader)
(c_i, c_j), (x_i_t, x_i_tk, x_i_prime,
x_j) = self.permute_data(data)
# encode
enc_i_t, enc_i_tk, enc_i_prime, enc_j = self.encode_step(
x_i_t, x_i_tk, x_i_prime, x_j)
# latent discriminate
latent_w_dis, latent_gp = self.latent_discriminate_step(
enc_i_t, enc_i_tk, enc_i_prime, enc_j)
lat_loss = -hps.alpha_dis * latent_w_dis + hps.lambda_ * latent_gp
reset_grad([self.LatentDiscriminator])
lat_loss.backward()
grad_clip([self.LatentDiscriminator],
self.hps.max_grad_norm)
self.lat_opt.step()
# print info
info = {
f'{flag}/D_latent_w_dis': latent_w_dis.data[0],
f'{flag}/latent_gp': latent_gp.data[0],
}
slot_value = (step, iteration + 1, hps.iters) + \
tuple([value for value in info.values()])
log = 'lat_D-%d:[%06d/%06d], w_dis=%.3f, gp=%.2f'
print(log % slot_value)
for tag, value in info.items():
self.logger.scalar_summary(tag, value, iteration)
# two stage training
if iteration >= hps.patch_start_iter:
for step in range(hps.n_patch_steps):
#===================== Train patch discriminator =====================#
data = next(self.data_loader)
(c_i, _), (x_i_t, _, _, _) = self.permute_data(data)
# encode
enc_i_t, = self.encode_step(x_i_t)
c_sample = self.sample_c(x_i_t.size(0))
x_tilde = self.decode_step(enc_i_t, c_i)
# Aux classify loss
patch_w_dis, real_logits, fake_logits, patch_gp = \
self.patch_discriminate_step(x_i_t, x_tilde, cal_gp=True)
patch_loss = -hps.beta_dis * patch_w_dis + hps.lambda_ * patch_gp + hps.beta_clf * c_loss
reset_grad([self.PatchDiscriminator])
patch_loss.backward()
grad_clip([self.PatchDiscriminator],
self.hps.max_grad_norm)
self.patch_opt.step()
# print info
info = {
f'{flag}/D_patch_w_dis': patch_w_dis.data[0],
f'{flag}/patch_gp': patch_gp.data[0],
f'{flag}/c_loss': c_loss.data[0],
f'{flag}/real_acc': real_acc,
f'{flag}/fake_acc': fake_acc,
}
slot_value = (step, iteration + 1, hps.iters) + \
tuple([value for value in info.values()])
log = 'patch_D-%d:[%06d/%06d], w_dis=%.3f, gp=%.2f, c_loss=%.3f, real_acc=%.2f, fake_acc=%.2f'
print(log % slot_value)
for tag, value in info.items():
self.logger.scalar_summary(tag, value, iteration)
#===================== Train G =====================#
data = next(self.data_loader)
(c_i, c_j), (x_i_t, x_i_tk, x_i_prime,
x_j) = self.permute_data(data)
# encode
enc_i_t, enc_i_tk, enc_i_prime, enc_j = self.encode_step(
x_i_t, x_i_tk, x_i_prime, x_j)
# decode
x_tilde = self.decode_step(enc_i_t, c_i)
loss_rec = torch.mean(torch.abs(x_tilde - x_i_t))
# latent discriminate
loss_adv = self.latent_discriminate_step(enc_i_t,
enc_i_tk,
enc_i_prime,
enc_j,
is_dis=False)
ae_loss = loss_rec + current_alpha * loss_adv
reset_grad([self.Encoder, self.Decoder])
retain_graph = True if hps.n_patch_steps > 0 else False
ae_loss.backward(retain_graph=retain_graph)
grad_clip([self.Encoder, self.Decoder], self.hps.max_grad_norm)
self.ae_opt.step()
info = {
f'{flag}/loss_rec': loss_rec.data[0],
f'{flag}/loss_adv': loss_adv.data[0],
f'{flag}/alpha': current_alpha,
}
slot_value = (iteration + 1, hps.iters) + tuple(
[value for value in info.values()])
log = 'G:[%06d/%06d], loss_rec=%.2f, loss_adv=%.2f, alpha=%.2e'
print(log % slot_value)
for tag, value in info.items():
self.logger.scalar_summary(tag, value, iteration + 1)
# patch discriminate
if hps.n_patch_steps > 0 and iteration >= hps.patch_start_iter:
c_sample = self.sample_c(x_i_t.size(0))
x_tilde = self.decode_step(enc_i_t, c_sample)
patch_w_dis, real_logits, fake_logits = \
self.patch_discriminate_step(x_i_t, x_tilde, cal_gp=False)
patch_loss = hps.beta_dec * patch_w_dis + hps.beta_clf * c_loss
reset_grad([self.Decoder])
patch_loss.backward()
grad_clip([self.Decoder], self.hps.max_grad_norm)
self.decoder_opt.step()
info = {
f'{flag}/G_patch_w_dis': patch_w_dis.data[0],
f'{flag}/c_loss': c_loss.data[0],
f'{flag}/real_acc': real_acc,
f'{flag}/fake_acc': fake_acc,
}
slot_value = (iteration + 1, hps.iters) + tuple(
[value for value in info.values()])
log = 'G:[%06d/%06d]: patch_w_dis=%.2f, c_loss=%.2f, real_acc=%.2f, fake_acc=%.2f'
print(log % slot_value)
for tag, value in info.items():
self.logger.scalar_summary(tag, value, iteration + 1)
if iteration % 1000 == 0 or iteration + 1 == hps.iters:
self.save_model(model_path, iteration)
def train(config):
train_data, word2index, tag2index, intent2index = preprocessing(
config.file_path, config.max_length)
if train_data == None:
print("Please check your data or its path")
return
encoder = Encoder(len(word2index), config.embedding_size,
config.hidden_size)
decoder = Decoder(len(tag2index), len(intent2index),
len(tag2index) // 3, config.hidden_size * 2)
if USE_CUDA:
encoder = encoder.cuda()
decoder = decoder.cuda()
#print("來到這裏了！1！")
encoder.init_weights()
decoder.init_weights()
#print("來到這裏了！2！")
loss_function_1 = nn.CrossEntropyLoss(ignore_index=0)
loss_function_2 = nn.CrossEntropyLoss()
enc_optim = optim.Adam(encoder.parameters(), lr=config.learning_rate)
dec_optim = optim.Adam(decoder.parameters(), lr=config.learning_rate)
#print("來到這裏了！3！")
for step in range(config.step_size):
losses = []
for i, batch in enumerate(getBatch(config.batch_size, train_data)):
x, y_1, y_2 = zip(*batch)
x = torch.cat(x)
tag_target = torch.cat(y_1)
intent_target = torch.cat(y_2)
# print("來到這裏了！4！")
x_mask = torch.cat([
Variable(torch.ByteTensor(tuple(map(lambda s: s == 0,
t.data)))).cuda()
if USE_CUDA else Variable(
torch.ByteTensor(tuple(map(lambda s: s == 0, t.data))))
for t in x
]).view(config.batch_size, -1)
y_1_mask = torch.cat([
Variable(torch.ByteTensor(tuple(map(lambda s: s == 0,
t.data)))).cuda()
if USE_CUDA else Variable(
torch.ByteTensor(tuple(map(lambda s: s == 0, t.data))))
for t in tag_target
]).view(config.batch_size, -1)
# print("來到這裏了！5！")
encoder.zero_grad()
decoder.zero_grad()
# print("來到這裏了！6！")
output, hidden_c = encoder(x, x_mask)
# print("來到這裏了！7！")
start_decode = Variable(
torch.LongTensor([
[word2index['<SOS>']] * config.batch_size
])).cuda().transpose(1, 0) if USE_CUDA else Variable(
torch.LongTensor([[word2index['<SOS>']] *
config.batch_size])).transpose(1, 0)
# print("來到這裏了！8！")
tag_score, intent_score = decoder(start_decode, hidden_c, output,
x_mask)
#print("來到這裏了！9！")
loss_1 = loss_function_1(tag_score, tag_target.view(-1))
# print("來到這裏了！10！")
loss_2 = loss_function_2(intent_score, intent_target)
#print("來到這裏了！11！")
loss = loss_1 + loss_2
losses.append(
loss.data.cpu().numpy() if USE_CUDA else loss.data.numpy())
#print("來到這裏了！12！")
loss.backward()
# print("來到這裏了！13！")
torch.nn.utils.clip_grad_norm(encoder.parameters(), 5.0)
torch.nn.utils.clip_grad_norm(decoder.parameters(), 5.0)
enc_optim.step()
dec_optim.step()
if i % 100 == 0:
with open("result.txt", "a+") as f:
#print("Step",step," epoch",i," : ",np.mean(losses))
print(f"Step是{step},epoch是{i} ：均值为{np.mean(losses)}")
f.write(f"Step是{step},epoch是{i} ：均值为{np.mean(losses)}")
f.write("\n")
losses = []
if not os.path.exists(config.model_dir):
os.makedirs(config.model_dir)
#print("來到這裏了！5！")
torch.save(decoder.state_dict(),
os.path.join(config.model_dir, 'jointnlu-decoder.pkl'))
torch.save(encoder.state_dict(),
os.path.join(config.model_dir, 'jointnlu-encoder.pkl'))
print("Train Complete!")
def train(config):
data_loader = DataLoader(config.file_path, config.max_length,
config.batch_size)
train_data, word2index, tag2index, intent2index = data_loader.load_train()
if train_data is None:
print("Please check your data or its path")
return
encoder = Encoder(len(word2index), config.embedding_size,
config.hidden_size)
decoder = Decoder(len(tag2index), len(intent2index),
config.hidden_size * 2)
if USE_CUDA:
encoder = encoder.cuda()
decoder = decoder.cuda()
encoder.init_weights()
decoder.init_weights()
loss_function_1 = nn.CrossEntropyLoss(ignore_index=0)
loss_function_2 = nn.CrossEntropyLoss()
enc_optim = optim.Adam(encoder.parameters(), lr=config.learning_rate)
dec_optim = optim.Adam(decoder.parameters(), lr=config.learning_rate)
for step in range(config.step_size):
losses = []
for i, batch in enumerate(data_loader.get_batch(train_data)):
x, embedding_x, y_1, y_2 = zip(*batch)
x = torch.cat(x)
embedding_x = torch.cat(embedding_x)
tag_target = torch.cat(y_1)
intent_target = torch.cat(y_2)
x_mask = torch.cat([
Variable(torch.ByteTensor(tuple(map(lambda s: s == 0,
t.data)))).cuda()
if USE_CUDA else Variable(
torch.ByteTensor(tuple(map(lambda s: s == 0, t.data))))
for t in x
]).view(len(batch), -1)
encoder.zero_grad()
decoder.zero_grad()
output, hidden_c = encoder(x, embedding_x, x_mask)
start_decode = Variable(
torch.LongTensor(
[[word2index['<SOS>']] * len(batch)])).cuda().transpose(
1, 0) if USE_CUDA else Variable(
torch.LongTensor([[word2index['<SOS>']] *
len(batch)])).transpose(1, 0)
tag_score, intent_score = decoder(start_decode, hidden_c, output,
x_mask)
loss_1 = loss_function_1(tag_score, tag_target.view(-1))
loss_2 = loss_function_2(intent_score, intent_target)
loss = loss_1 + loss_2
losses.append(
loss.data.cpu().numpy() if USE_CUDA else loss.data.numpy())
loss.backward()
torch.nn.utils.clip_grad_norm(encoder.parameters(), 5.0)
torch.nn.utils.clip_grad_norm(decoder.parameters(), 5.0)
enc_optim.step()
dec_optim.step()
if i % 100 == 0:
print("Step", step, " : ", np.mean(losses))
losses = []
if not os.path.exists(config.model_dir):
os.makedirs(config.model_dir)
torch.save(encoder, os.path.join(config.model_dir, 'jointnlu-encoder.pt'))
torch.save(decoder, os.path.join(config.model_dir, 'jointnlu-decoder.pt'))
print("Training Complete!")
class sample:
def __init__(self):
parser = argparse.ArgumentParser(description='Image Captioning')
parser.add_argument('--root',
default='../../../cocodataset/',
type=str)
parser.add_argument(
'--sample_image',
default='../../../cocodataset/val2017/000000435205.jpg',
type=str)
parser.add_argument('--epochs', default=100, type=int)
parser.add_argument('--lr', default=1e-4, type=float)
parser.add_argument('--batch_size', default=128, help='')
parser.add_argument('--num_workers', default=4, type=int)
parser.add_argument('--embed_dim', default=256, type=int)
parser.add_argument('--hidden_size', default=512, type=int)
parser.add_argument('--num_layers', default=1, type=int)
parser.add_argument('--encoder_path',
default='./model/Encoder-100.ckpt',
type=str)
parser.add_argument('--decoder_path',
default='./model/Decoder-100.ckpt',
type=str)
parser.add_argument('--vocab_path', default='./vocab/', type=str)
self.args = parser.parse_args()
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, ), (0.5, )),
transforms.Resize((224, 224))
])
with open(self.args.vocab_path + 'vocab.pickle', 'rb') as f:
data = pickle.load(f)
self.vocab = data
self.DataLoader = get_dataloader(root=self.args.root,
transform=self.transform,
shuffle=True,
batch_size=self.args.batch_size,
num_workers=self.args.num_workers,
vocab=self.vocab)
self.Encoder = Encoder(embed_dim=self.args.embed_dim)
self.Decoder = Decoder(embed_dim=self.args.embed_dim,
hidden_size=self.args.hidden_size,
vocab_size=len(self.vocab),
num_layers=self.args.num_layers)
def load_image(self, image_path):
image = Image.open(image_path).convert('RGB')
image = self.transform(image).unsqueeze(0)
return image
def main(self):
self.Encoder.load_state_dict(torch.load(self.args.encoder_path))
self.Decoder.load_state_dict(torch.load(self.args.decoder_path))
self.Encoder = self.Encoder.cuda().eval()
self.Decoder = self.Decoder.cuda().eval()
sample_image = self.load_image(self.args.sample_image).cuda()
output = self.Encoder(sample_image)
output = self.Decoder.sample(output)[0].cpu().numpy()
sample_caption = []
for idx in output:
word = self.vocab.idx2word[idx]
sample_caption.append(word)
if word == '<end>':
break
sentence = ' '.join(sample_caption)
print(sentence)
def train(config, encoder_in = None, decoder_in = None):
train_data, word2index, tag2index, intent2index = preprocessing(config.file_path,config.max_length)
if train_data==None:
print("Please check your data or its path")
return
if encoder_in != None:
encoder = encoder_in
decoder = decoder_in
else:
encoder = Encoder(len(word2index),config.embedding_size,config.hidden_size)
decoder = Decoder(len(tag2index),len(intent2index),len(tag2index)//3,config.hidden_size*2)
if USE_CUDA:
encoder = encoder.cuda()
decoder = decoder.cuda()
encoder.init_weights()
decoder.init_weights()
loss_function_1 = nn.CrossEntropyLoss(ignore_index=0)
loss_function_2 = nn.CrossEntropyLoss()
enc_optim= optim.Adam(encoder.parameters(), lr=config.learning_rate)
dec_optim = optim.Adam(decoder.parameters(),lr=config.learning_rate)
for step in range(config.step_size):
losses=[]
for i, batch in enumerate(getBatch(config.batch_size,train_data)):
x,y_1,y_2 = zip(*batch) # sin,sout,intent
x = torch.cat(x)
tag_target = torch.cat(y_1)
intent_target = torch.cat(y_2)
x_mask = torch.cat([Variable(torch.ByteTensor(tuple(map(lambda s: s ==0, t.data)))).cuda() if USE_CUDA else Variable(torch.ByteTensor(tuple(map(lambda s: s ==0, t.data)))) for t in x]).view(config.batch_size,-1)
y_1_mask = torch.cat([Variable(torch.ByteTensor(tuple(map(lambda s: s ==0, t.data)))).cuda() if USE_CUDA else Variable(torch.ByteTensor(tuple(map(lambda s: s ==0, t.data)))) for t in tag_target]).view(config.batch_size,-1)
encoder.zero_grad()
decoder.zero_grad()
output, hidden_c = encoder(x,x_mask)
start_decode = Variable(torch.LongTensor([[word2index['<SOS>']]*config.batch_size])).cuda().transpose(1,0) if USE_CUDA else Variable(torch.LongTensor([[word2index['<SOS>']]*config.batch_size])).transpose(1,0)
tag_score, intent_score = decoder(start_decode,hidden_c,output,x_mask)
loss_1 = loss_function_1(tag_score,tag_target.view(-1))
loss_2 = loss_function_2(intent_score,intent_target)
loss = loss_1+loss_2
losses.append(loss.data.cpu().numpy() if USE_CUDA else loss.data.numpy())
loss.backward()
torch.nn.utils.clip_grad_norm(encoder.parameters(), 5.0)
torch.nn.utils.clip_grad_norm(decoder.parameters(), 5.0)
enc_optim.step()
dec_optim.step()
if i % 100==0:
print("Step",step," epoch",i," : ",np.mean(losses))
losses=[]
t = Check()
t.test(encoder,decoder)
count = t.test_error_count
rate = t.test_error_rate
if not os.path.exists(config.model_dir):
os.makedirs(config.model_dir)
torch.save(decoder, os.path.join(config.model_dir, str(count)+'_'+str(rate)+'_'+'decoder.pkl'))
torch.save(encoder, os.path.join(config.model_dir, str(count)+'_'+str(rate)+'_'+'encoder.pkl'))
# if not os.path.exists(config.model_dir):
# os.makedirs(config.model_dir)
# torch.save(decoder.state_dict(),os.path.join(config.model_dir,'jointnlu-decoder.pkl'))
# torch.save(encoder.state_dict(),os.path.join(config.model_dir, 'jointnlu-encoder.pkl'))
# torch.save(decoder,os.path.join(config.model_dir,'jointnlu-decoder.pkl'))
# torch.save(encoder,os.path.join(config.model_dir, 'jointnlu-encoder.pkl'))
print("Train Complete!")
with open("../data/vocab.pkl", 'rb') as f:
vocab = pickle.load(f)
dataloader = get_loader("../data/resized/",
"../data/annotations/captions_train2014.json",
vocab,
trans,
128,
shuffle=True)
encoder = Encoder(256)
decoder = Decoder(256, 512, len(vocab), 1)
if torch.cuda.is_available():
encoder.cuda()
decoder.cuda()
criterion = nn.CrossEntropyLoss()
params = list(decoder.parameters()) + list(encoder.parameters()) + list(
encoder.bn.parameters())
optimizer = torch.optim.Adam(params, lr=0.001)
total_step = len(dataloader)
for epoch in range(5):
for i, (images, captions, lengths) in enumerate(dataloader):
images = to_var(images, volatile=True)
captions = to_var(captions)
targets = pack_padded_sequence(captions, lengths,
batch_first=True)[0]
decoder.zero_grad()
encoder.zero_grad()
class TotalSytle():
def __init__(self):
self.train_loader = get_data_loader(content_path=CONTENT_PATH,
style_path=STYLE_PATH,
batch_size=BATCH_SIZE,
small_test=False)
self.encoder = Encoder()
self.decoder = Decoder()
self.decoder.load_state_dict(torch.load("weights/zf_decoder.pt"))
self.mse_loss = torch.nn.MSELoss()
# self.style_loss = StyleLoss() ## TODO: Complete Styleloss
# self.content_loss = ContentLoss() ## TODO: Complete ContentLoss
print(
"----------------------Model is loaded----------------------------"
)
parameters = self.decoder.parameters()
self.optimizer = torch.optim.Adam(parameters, lr=LEARNING_RATE)
self.use_gpu = True
if self.use_gpu:
print(
"----------------------GPU is used to train---------------------------"
)
self.encoder.cuda()
self.decoder.cuda()
self.mse_loss.cuda()
else:
print(
"----------------------CPU is used to train----------------------------"
)
self.alpha = 0.5 # the weight of content loss and style loss
self.beta = 1 # the weight of inter_scale loss and inter_scale loss. 1: only use loss_intra_scale
def train(self):
for epoch in range(MAX_EPOCH):
total_loss = 0
for batch_id, (style_imgs,
content_imgs) in enumerate(self.train_loader):
style_imgs = style_imgs.cuda()
content_imgs = content_imgs.cuda()
self.optimizer.zero_grad()
torch.cuda.empty_cache()
# Parse the style_imgs and content_imgs into encoder
encoded_style, output_style = self.encoder(style_imgs)
encoded_content, output_content = self.encoder(content_imgs)
# e_loss = encoded_content[-1]
encoded_content_save = copy.deepcopy(encoded_content)
encoded_style_save = copy.deepcopy(encoded_style)
# Compute the MST transformed relu
relu1_2, relu2_2, relu3_3 = MST(encoded_content, encoded_style)
# print(relu1_2)
# print(relu1_2.size(), relu2_2.size(), relu3_3.size())
# Skip connection with decoder
stylized_img = self.decoder(relu1_2, relu2_2, relu3_3)
# print(stylized_img.size())
# Extract the features of generated stylized img
encoded_stylized, _ = self.encoder(stylized_img)
# content_loss, _ = self.encoder(content_imgs)
# compute the loss between stylized imgs and content imgs
# use only relu3_3 to as the 'content' of an img
loss_c = self.mse_loss(encoded_stylized[-1],
encoded_content_save[-1])
# loss_c = calc_style_loss(encoded_stylized[-1], encoded_content_save[-1], self.mse_loss)
# compute the loss between stylized imgs and style imgs
# intra scale loss
loss_intra_scale = calc_style_loss(encoded_stylized[0],
encoded_style_save[0],
self.mse_loss)
# loss_intra_scale = self.mse_loss(encoded_stylized[0], encoded_style_save[0])
for i in range(1, 3):
loss_intra_scale += calc_style_loss(
encoded_stylized[i], encoded_style_save[i],
self.mse_loss)
#loss_intra_scale += self.mse_loss(encoded_stylized[i], encoded_style_save[i])
# inter scale loss
encoded_stylized = upsample_and_cat(encoded_stylized)
encoded_content = upsample_and_cat(encoded_content)
loss_inter_sacle = calc_style_loss(encoded_stylized,
encoded_content,
self.mse_loss)
# weighted sum of inter_scale loss and intra scale loss
# the default self.bata = 1 for only using intra scale loss
loss_s = self.beta * loss_intra_scale + \
(1-self.beta) * loss_inter_sacle
# weighted sum of style loss and content loss
# loss = self.alpha * loss_s + (1-self.alpha) * loss_c
loss = 0.99 * loss_s + 0.01 * loss_c
# print("loss_s-loss_c: ", loss_s.item(), loss_c.item())
# print(loss.item())
loss.backward()
total_loss += loss.item()
self.optimizer.step()
generated_img = stylized_img.detach().cpu()
for i, gimg in enumerate(generated_img):
# gimg = transforms.functional.to_pil_image(gimg)
vutils.save_image(
gimg,
"./data/generate/" + str(batch_id) + str(i) + ".jpg")
# gimg.save("./data/generate/" + str(batch_id) + str(i) + ".jpg")
# generated_img = stylized_img.detach().cpu()
# generated_img = transforms.functional.to_pil_image(generated_img[0])
# vutils.save_image(generated_img,"./data/generate/" + str(epoch) + ".jpg")
# generated_img.save("./data/generate/" + str(epoch) + ".jpg")
print("[TRAIN] EPOCH %d/%d, Loss/batch_num: %.4f" %
(epoch, MAX_EPOCH, total_loss / (batch_id + 1)))
if epoch % 50 == 0:
torch.save(self.encoder.state_dict(),
"./weights/epoch" + str(epoch) + "_encoder.pt")
torch.save(self.decoder.state_dict(),
"./weights/epoch" + str(epoch) + "_decoder.pt")