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 build_model(self, wavenet_mel):
hps = self.hps
ns = self.hps.ns
emb_size = self.hps.emb_size
c = 80 if wavenet_mel else 513
patch_classify_kernel = (3, 4) if wavenet_mel else (17, 4)
self.Encoder = cc(Encoder(c_in=c, ns=ns, dp=hps.enc_dp))
self.Decoder = cc(
Decoder(c_out=c, ns=ns, c_a=hps.n_speakers, emb_size=emb_size))
self.Generator = cc(
Decoder(c_out=c, ns=ns, c_a=hps.n_speakers, emb_size=emb_size))
self.SpeakerClassifier = cc(
SpeakerClassifier(ns=ns, n_class=hps.n_speakers, dp=hps.dis_dp))
self.PatchDiscriminator = cc(
nn.DataParallel(
PatchDiscriminator(
ns=ns,
n_class=hps.n_speakers,
classify_kernel_size=patch_classify_kernel)))
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.clf_opt = optim.Adam(self.SpeakerClassifier.parameters(),
lr=self.hps.lr,
betas=betas)
self.gen_opt = optim.Adam(self.Generator.parameters(),
lr=self.hps.lr,
betas=betas)
self.patch_opt = optim.Adam(self.PatchDiscriminator.parameters(),
lr=self.hps.lr,
betas=betas)
def build_model(self):
hps = self.hps
ns = self.hps.ns
emb_size = self.hps.emb_size
self.Encoder = cc(Encoder(ns=ns, dp=hps.enc_dp))
self.Decoder = cc(Decoder(ns=ns, c_a=hps.n_speakers,
emb_size=emb_size))
self.Generator = cc(
Decoder(ns=ns, c_a=hps.n_speakers, emb_size=emb_size))
self.SpeakerClassifier = cc(
SpeakerClassifier(ns=ns, n_class=hps.n_speakers, dp=hps.dis_dp))
self.PatchDiscriminator = cc(
nn.DataParallel(PatchDiscriminator(ns=ns, n_class=hps.n_speakers)))
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.clf_opt = optim.Adam(self.SpeakerClassifier.parameters(),
lr=self.hps.lr,
betas=betas)
self.gen_opt = optim.Adam(self.Generator.parameters(),
lr=self.hps.lr,
betas=betas)
self.patch_opt = optim.Adam(self.PatchDiscriminator.parameters(),
lr=self.hps.lr,
betas=betas)
def __init__(self, src_vocab, tgt_vocab,
max_len=300, hidden_size=300, n_layers=2, clip=5, n_epochs=30):
# hyper-parameters
self.max_len = max_len
self.hidden_size = hidden_size
self.n_layers = n_layers
self.clip = clip
self.n_epochs = n_epochs
# vocab
self.src_vocab = src_vocab
self.tgt_vocab = tgt_vocab
self.pad_idx = self.src_vocab.stoi[PAD]
# prepare model
self.encoder = Encoder(self.src_vocab, self.max_len, self.hidden_size, self.n_layers)
self.decoder = Decoder(self.tgt_vocab, self.max_len, self.hidden_size * 2, self.n_layers)
self.reverse_decoder = Decoder(self.tgt_vocab, self.max_len, self.hidden_size * 2, self.n_layers, reverse=True)
self.model = Seq2SeqConcat(self.encoder, self.decoder, self.reverse_decoder, self.pad_idx)
self.model.to(device)
print(self.model)
print("Total parameters:", sum([p.nelement() for p in self.model.parameters()]))
# initialize weights
for name, param in self.model.named_parameters():
if "lstm.bias" in name:
# set lstm forget gate to 1 (Jozefowicz et al., 2015)
n = param.size(0)
param.data[n//4:n//2].fill_(1.0)
elif "lstm.weight" in name:
nn.init.xavier_uniform_(param)
# prepare loss function; don't calculate loss on PAD tokens
self.criterion = nn.NLLLoss(ignore_index=self.pad_idx)
# prepare optimizer and scheduler
self.optimizer = Adam(self.model.parameters())
self.scheduler = CyclicLR(self.optimizer, base_lr=0.00001, max_lr=0.00005,
step_size_up=4000, step_size_down=4000,
mode="triangular", gamma=1.0, cycle_momentum=False)
# book keeping vars
self.global_iter = 0
self.global_numel = []
self.global_loss = []
self.global_acc = []
# visualization
self.vis_loss = Visualization(env_name="aivivn_tone", xlabel="step", ylabel="loss", title="loss (mean per 300 steps)")
self.vis_acc = Visualization(env_name="aivivn_tone", xlabel="step", ylabel="acc", title="training accuracy (mean per 300 steps)")
def create_model(args):
trim_min_count = 5
data_loader = DataLoader(args, trim_min_count=trim_min_count)
embed_model = nn.Embedding(data_loader.vocab_len, args.embed)
embed_model.weight.data.copy_(data_loader.embed_vectors)
encode_model = Encoder(
embed_model=embed_model,
hidden_size=args.hidden,
span_size=data_loader.span_size,
dropout=args.dropout,
)
decode_model = Decoder(
embed_model=embed_model,
op_set=data_loader.op_set,
vocab_dict=data_loader.vocab_dict,
class_list=data_loader.class_list,
hidden_size=args.hidden,
dropout=args.dropout,
use_cuda=args.use_cuda
)
seq2seq = Seq2seq(encode_model, decode_model)
return seq2seq, data_loader
def inference(speaker, text):
BATCH_SIZE = 1
encoder = Encoder(BATCH_SIZE, len(vocabulary.VOCAB), len(dataset.SPEAKERS))
decoder = Decoder(BATCH_SIZE, hparams.n_mel * 4, (hparams.n_fft // 2 + 1) * 4)
global_step = tf.Variable(0, dtype=tf.int64) # summary 의 step 은 int64
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
checkpoint_dir = './checkpoint'
checkpoint = tf.train.Checkpoint(optimizer=optimizer,
encoder=encoder,
decoder=decoder,
global_step=global_step)
checkpoint_manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=5)
checkpoint.restore(checkpoint_manager.latest_checkpoint)
if checkpoint_manager.latest_checkpoint:
print("체크포인트 {} 에서 로드합니다...".format(checkpoint_manager.latest_checkpoint))
print(global_step.numpy(), "step 에서 추론을 진행합니다...")
tokens = np.array([vocabulary.add_eos(vocabulary.str2token(text))])
speaker = np.array([dataset.SPEAKER_NAME2IDX[speaker]])
key, value = encoder(tokens, speaker)
dec_out, weights, linear = decoder.call_inference(key, value)
return dec_out[0].numpy().reshape(-1, 80), weights[0].numpy(), linear[0].numpy().reshape(-1, hparams.n_fft // 2 + 1)
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 build_model(config, input_vocab_size, target_vocab_size):
embed_size = config['embed_size']
hidden_size = config['hidden_size']
proj_size = config['proj_size']
enc_num_layers = config['enc_num_layers']
dec_num_layers = config['dec_num_layers']
dropout = config['dropout']
attn_type = config['attn_type']
self_attn = config['self_attn']
intra_temp_attn = config['intra_temp_attn']
dec_attn = config['dec_attn']
if self_attn or intra_temp_attn:
dec_attn = True
encoder = Encoder(input_vocab_size, embed_size, hidden_size,
enc_num_layers, dropout)
decoder = Decoder(target_vocab_size,
embed_size,
hidden_size,
dec_num_layers,
proj_size,
dropout,
attn_type=attn_type,
self_attn=self_attn,
dec_attn=dec_attn,
intra_temp_attn=intra_temp_attn)
model = Seq2Seq(encoder, decoder).to(device)
return model
def test():
if (test_content_image == '' or test_style_image == ''):
print("Please input the content and style image!")
return
# 设置设备为GPU/cpu
if torch.cuda.is_available() and gpu >= 0:
device = torch.device(f'cuda:{gpu}')
print(f'# CUDA available: {torch.cuda.get_device_name(0)}')
else:
device = 'cpu'
# device = 'cpu'
# 载入编码器和解码器
encoder = VGGEncoder().to(device)
decoder = Decoder()
decoder.load_state_dict(torch.load(model_state_path))
decoder = decoder.to(device)
try:
content = Image.open(test_content_image)
style = Image.open(test_style_image)
c_tensor = trans(content).unsqueeze(0).to(device)
s_tensor = trans(style).unsqueeze(0).to(device)
# 不记录torch自带的计算图
with torch.no_grad():
cf = encoder(c_tensor)
sf = encoder(s_tensor)
style_swap_res = style_swap(cf, sf, patch_size, 1)
del cf
del sf
del encoder
out = decoder(style_swap_res)
c_denorm = denorm(c_tensor, device)
out_denorm = denorm(out, device)
res = torch.cat([c_denorm, out_denorm], dim=0)
res = res.to('cpu')
except RuntimeError as e:
traceback.print_exc()
print('Images are too large to transfer.')
if default_output_name == '':
c_name = os.path.splitext(os.path.basename(test_content_image))[0]
s_name = os.path.splitext(os.path.basename(test_style_image))[0]
output_name = f'{c_name}_{s_name}'
else:
output_name = default_output_name
try:
save_image(out_denorm, f'{output_name}.jpg')
save_image(res, f'{output_name}_pair.jpg', nrow=2)
o = Image.open(f'{output_name}_pair.jpg')
style = style.resize((i // 4 for i in content.size))
box = (o.width // 2, o.height - style.height)
o.paste(style, box)
o.save(f'{output_name}_style_transfer_demo.jpg', quality=95)
print(f'result saved into files starting with {output_name}')
except:
pass
def transform(sentence, l1, l2, op_lang):
if op_lang == "hi":
dir = "training_checkpoints"
with open("tensors1.pkl", 'rb') as f:
example_input_batch = pickle.load(f)
example_target_batch = pickle.load(f)
elif op_lang == "nep":
dir = "training_nepali"
with open("tensors_nep.pkl", 'rb') as f:
example_input_batch = pickle.load(f)
example_target_batch = pickle.load(f)
elif op_lang == "pun":
dir = "training_punjabi"
with open("tensors_pun.pkl", 'rb') as f:
example_input_batch = pickle.load(f)
example_target_batch = pickle.load(f)
print(example_input_batch)
global input_lang, output_lang
input_lang = l1
output_lang = l2
global BATCH_SIZE, units, embedding_dim
BATCH_SIZE = 64
units = 512
if op_lang == "pun":
units = 1024
embedding_dim = 256
vocab_inp_size = len(input_lang.word2index) + 1
vocab_tar_size = len(output_lang.word2index) + 1
global encoder, decoder, optimizer
optimizer = tf.keras.optimizers.Adam()
loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction='none')
encoder = Encoder(vocab_inp_size, embedding_dim, units, BATCH_SIZE)
sample_hidden = encoder.initialize_hidden_state()
sample_output, sample_hidden = encoder(example_input_batch, sample_hidden)
attention_layer = BahdanauAttention(10)
attention_result, attention_weights = attention_layer(
sample_hidden, sample_output)
decoder = Decoder(vocab_tar_size, embedding_dim, units, BATCH_SIZE)
sample_decoder_output, _, _ = decoder(tf.random.uniform((BATCH_SIZE, 1)),
sample_hidden, sample_output)
checkpoint = tf.train.Checkpoint(optimizer=optimizer,
encoder=encoder,
decoder=decoder)
curr = pathlib.Path(__file__)
print(encoder.summary())
temp = os.path.join(curr.parent, "available_models")
checkpoint_dir = os.path.join(temp, dir)
print(checkpoint_dir)
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
print(checkpoint_dir)
print(checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir)))
# print(checkpoint.restore(checkpoint_dir))
print(encoder.summary())
return translate(sentence)
def init_model(ixtoword):
if cfg.CAP.USE_ORIGINAL:
caption_cnn = CAPTION_CNN(embed_size=cfg.CAP.EMBED_SIZE)
caption_rnn = CAPTION_RNN(embed_size=cfg.CAP.EMBED_SIZE,
hidden_size=cfg.CAP.HIDDEN_SIZE,
vocab_size=len(ixtoword),
num_layers=cfg.CAP.NUM_LAYERS)
else:
caption_cnn = Encoder()
caption_rnn = Decoder(idx2word=ixtoword)
decoder_optimizer = torch.optim.Adam(params=caption_rnn.parameters(),
lr=cfg.CAP.LEARNING_RATE)
if cfg.CAP.CAPTION_CNN_PATH and cfg.CAP.CAPTION_RNN_PATH:
print('Pre-Trained Caption Model')
caption_cnn_checkpoint = torch.load(
cfg.CAP.CAPTION_CNN_PATH,
map_location=lambda storage, loc: storage)
caption_rnn_checkpoint = torch.load(
cfg.CAP.CAPTION_RNN_PATH,
map_location=lambda storage, loc: storage)
caption_cnn.load_state_dict(caption_cnn_checkpoint['model_state_dict'])
caption_rnn.load_state_dict(caption_rnn_checkpoint['model_state_dict'])
decoder_optimizer.load_state_dict(
caption_rnn_checkpoint['optimizer_state_dict'])
caption_cnn = caption_cnn.to(cfg.DEVICE)
caption_rnn = caption_rnn.to(cfg.DEVICE)
decoder_optimizer = decoder_optimizer
return caption_cnn, caption_rnn, decoder_optimizer
def model_fn(model_dir):
print("Loading model.")
model_info = {}
model_info_path = os.path.join(model_dir, 'model_info.pth')
with open(model_info_path, 'rb') as f:
model_info = torch.load(f)
print("model_info: {}".format(model_info))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
encoder = Encoder(model_info['input_dim'], model_info['hidden_dim'])
decoder = Decoder(model_info['input_dim'], model_info['output_dim'], model_info['hidden_dim'])
model = WordOrderer(encoder, decoder).to(device)
model_path = os.path.join(model_dir, 'model.pth')
with open(model_path, 'rb') as f:
model.load_state_dict(torch.load(f))
model_path = os.path.join(model_dir, 'model.pth')
with open(model_path, 'rb') as f:
model.load_state_dict(torch.load(f))
# Load the 2 saved letter2int_dict.
letter2int_dict_path = os.path.join(model_dir, 'letter2int_dict.pkl')
with open(letter2int_dict_path, 'rb') as f:
model.letter2int_dict = pickle.load(f)
int2letter_dict_path = os.path.join(model_dir, 'int2letter_dict.pkl')
with open(int2letter_dict_path, 'rb') as f:
model.int2letter_dict = pickle.load(f)
model.to(device).eval()
print("Done loading model.")
return model
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)
def da_rnn(train_data: TrainData, n_targs: int,
encoder_hidden_size=64, decoder_hidden_size=64,
T=10, learning_rate=0.01, batch_size=128):
train_cfg = TrainConfig(T, int(train_data.feats.shape[0] * 0.7), batch_size, nn.MSELoss())
logger.info(f"Training size: {train_cfg.train_size:d}.")
enc_kwargs = {"input_size": train_data.feats.shape[1], "hidden_size": encoder_hidden_size, "T": T}
encoder = Encoder(**enc_kwargs).to(device)
with open(os.path.join("data", "enc_kwargs.json"), "w") as f:
json.dump(enc_kwargs, f, indent=4)
dec_kwargs = {"encoder_hidden_size": encoder_hidden_size,
"decoder_hidden_size": decoder_hidden_size, "T": T, "out_feats": n_targs}
decoder = Decoder(**dec_kwargs).to(device)
with open(os.path.join("data", "dec_kwargs.json"), "w") as f:
json.dump(dec_kwargs, f, indent=4)
encoder_optimizer = optim.Adam(
params=[p for p in encoder.parameters() if p.requires_grad],
lr=learning_rate)
decoder_optimizer = optim.Adam(
params=[p for p in decoder.parameters() if p.requires_grad],
lr=learning_rate)
da_rnn_net = DaRnnNet(encoder, decoder, encoder_optimizer, decoder_optimizer)
return train_cfg, da_rnn_net
def __init__(self,
train_folder: str,
test_folder: str,
alpha: float = 0.01,
beta: float = 1.0,
predicted_poses: int = 20,
previous_poses: int = 10,
stride: int = None,
batch_size: int = 50,
embedding: str = None,
text_folder: str = None,
*args,
**kwargs):
super().__init__()
self.save_hyperparameters()
self.encoder = ContextEncoder(26, 150, 1)
self.decoder = Decoder(45, 150, 300, max_gen=predicted_poses)
self.predicted_poses = predicted_poses
self.previous_poses = previous_poses
self.loss = MSELoss()
self.train_folder = train_folder
self.test_folder = test_folder
self.alpha = alpha
self.beta = beta
self.stride = predicted_poses if stride is None else stride
self.batch_size = batch_size
if embedding is not None:
self.vocab = Vocab(embedding)
self.embedder = nn.Embedding(len(self.vocab.token_to_idx),
len(self.vocab.weights[0]),
_weight=torch.FloatTensor(
self.vocab.weights))
else:
self.vocab = None
self.text_folder = text_folder
def get_model(config):
encoder = maybe_cuda(Encoder(config), cuda=config.cuda)
if config.attn_method != "disabled":
decoder = maybe_cuda(AttnDecoder(config), cuda=config.cuda)
else:
decoder = maybe_cuda(Decoder(config), cuda=config.cuda)
return encoder, decoder
def main(args):
train_doc,train_summ,test_doc, \
test_summ,val_doc,val_summ=read_from_file(args.data_file)
vocab = read_vocab(args.path_to_vocab)
embeddings = Embeddings(args.embed_size, args.vocab_size).cuda()
encoder = Encoder(args.embed_size, args.hidden_size).cuda()
decoder = Decoder(args.embed_size, args.hidden_size,
args.vocab_size).cuda()
generator = BeamSearch(vocab, args.max_decode_len, args.min_decode_len,
args.beam_size)
trainloader = DataLoader(train_doc, train_summ, vocab, args.batch_size,
args.max_doc_len, args.max_summ_len)
testloader = Test_loader(test_doc, test_summ, vocab, args.max_doc_len,
args.max_summ_test_len)
valloader = Test_loader(val_doc, val_summ, vocab, args.max_doc_len,
args.max_summ_test_len)
if args.use_pretrained:
params = torch.load(args.pretrained_model)
embeddings.load_state_dict(params['embed_params'])
encoder.load_state_dict(params['encoder_params'])
decoder.load_state_dict(params['decoder_params'])
test(embeddings, encoder, decoder, testloader, generator, args.lambda_)
train(embeddings, encoder, decoder, generator, trainloader, valloader,
args.iterations, args.lambda_, args.lr, args.max_grad_norm,
args.initial_accum_val, args.threshold)
test(embeddings, encoder, decoder, testloader, generator, args.lambda_)
def run_func():
conf = json.load(open('config.json'))
print(conf)
if conf['data_name'] == "NYT":
config = NYT_Config(conf)
else:
config = WebNLG_Config(conf)
train = dataset(config.question_train, config.context_train,
config.answer_train, config.cnn_output_train,
config.cnn_list_train)
dev = dataset(config.question_dev, config.context_dev, config.answer_dev,
config.cnn_output_dev, config.cnn_list_dev)
test = dataset(config.question_test, config.context_test,
config.answer_test, config.cnn_output_test,
config.cnn_list_test)
print(len(train))
print(len(dev))
print(len(test))
encoder = Encoder(config.hidden_state_size)
decoder = Decoder(config.hidden_state_size)
cnn = CNN(config, is_training=True)
qa = QASystem(encoder, decoder, cnn, config)
sess = tf.Session()
qa.initialize_model(sess, config.train_dir)
qa.train(sess, [train, dev, test], config.train_dir, config)
def train(dct_size, embed_size=256, hidden_size=512, epochs=10, num_layers=1, save_step=1000, lr=0.001, model_save='model/'):
encoder = Encoder(embed_size=embed_size).to(device)
decoder = Decoder(embed_size=embed_size, hidden_size=hidden_size, dct_size=len(dct), num_layers=num_layers).to(device)
criterion = nn.CrossEntropyLoss()
params = list(decoder.parameters()) + list(encoder.linear.parameters()) + list(encoder.bn.parameters())
optimizer = torch.optim.Adam(params, lr=lr)
for epoch in range(epochs):
print(f'epoch {epoch+1}/{epochs}: ')
for i, (images, captions, lengths) in enumerate(tqdm(data_loader)):
# Set mini-batch dataset
images = images.to(device)
captions = captions.to(device)
targets = pack_padded_sequence(captions, lengths, batch_first=True)[0]
# Forward, backward and optimize
features = encoder(images)
outputs = decoder(features, captions, lengths)
loss = criterion(outputs, targets)
decoder.zero_grad()
encoder.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % save_step == 0:
torch.save(decoder.state_dict(), os.path.join(
model_save, 'decoder-{}-{}.ckpt'.format(epoch+1, i+1)))
torch.save(encoder.state_dict(), os.path.join(
model_save, 'encoder-{}-{}.ckpt'.format(epoch+1, i+1)))
def main(fpath):
ENC_EMB_DIM = 256
DEC_EMB_DIM = 256
ENC_HID_DIM = 512
DEC_HID_DIM = 512
ENC_DROPOUT = 0.5
DEC_DROPOUT = 0.5
device = torch.device('cuda')
dataset = Dataset()
INPUT_DIM = len(dataset.SRC.vocab)
OUTPUT_DIM = len(dataset.TRG.vocab)
SRC_PAD_IDX = dataset.SRC.vocab.stoi[dataset.SRC.pad_token]
encoder = Encoder(INPUT_DIM, ENC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM,
ENC_DROPOUT)
attention = Attention(ENC_HID_DIM, DEC_HID_DIM)
decoder = Decoder(DEC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, OUTPUT_DIM,
DEC_DROPOUT, attention)
model = Seq2Seq(encoder, decoder, SRC_PAD_IDX, device)
model.load_state_dict(torch.load("best_model.pt"))
model.to(device)
with open(fpath, "r") as f:
sentences = f.readlines()
translate_sentence(model, sentences, dataset.SRC, dataset.TRG, device)
def forward_propogate(INPUT_IMAGE, weights, biases):
mean, std = Encoder(INPUT_IMAGE).encode(weights, biases)
latent_layer = latent_space(mean, std)
decoder_output = Decoder(latent_layer).decode(weights, biases)
return decoder_output, mean, std
def model_setting(args):
# Encoder
encoder_model = EncoderResNet()
encoder_model = nn.DataParallel(encoder_model)
encoder_model.load_state_dict(torch.load(args.encoder_model_path))
# Decoder
decoder_model = Decoder(vis_dim=args.vis_dim,
vis_num=args.vis_num,
embed_dim=args.embed_dim,
hidden_dim=args.hidden_dim,
vocab_size=args.vocab_size,
num_layers=args.num_layers,
dropout_ratio=args.dropout_ratio)
decoder_model = nn.DataParallel(decoder_model)
decoder_model.load_state_dict(torch.load(args.decoder_model_path))
# move to GPU and change evaluation mode
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("DEVICE: {}".format(device))
encoder_model = encoder_model.to(device)
decoder_model = decoder_model.to(device)
return encoder_model, decoder_model
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_by_human():
# test the result from .de to .en
args = parse_arguments()
hidden_size = 512
embed_size = 256
assert torch.cuda.is_available()
print("[!] preparing dataset for test ...")
train_iter, val_iter, test_iter, DE, EN = load_dataset(args.batch_size)
de_size, en_size = len(DE.vocab), len(EN.vocab)
# load the model
encoder = Encoder(de_size,
embed_size,
hidden_size,
n_layers=2,
dropout=0.5)
decoder = Decoder(embed_size,
hidden_size,
en_size,
n_layers=1,
dropout=0.5)
seq2seq = Seq2Seq(encoder, decoder).cuda()
seq2seq.load_state_dict(torch.load('.save/seq2seq_21.pt'))
# only decoder 1 batch sents
for b, batch in enumerate(train_iter):
src, len_src = batch.src
trg, len_trg = batch.trg
src, trg = src.cuda(), trg.cuda()
# do not use force teaching, just use the maximum possibility
output = seq2seq(src, trg, 0)
output = output.transpose(0, 1) # (B*T*N)
src = src.transpose(0, 1) # (B*T)
# src
for source, result in zip(src, output):
print('German: : ')
print(' ', end=' ')
for word in source:
if DE.vocab.itos[word] in ["<pad>", "<sos>", "<unk>", "<eos>"]:
continue
print(DE.vocab.itos[word], end=' ')
# print(word, end=' ')
print()
print('English: ')
print(' ', end=' ')
for word in result:
_, index = word.max(0)
if EN.vocab.itos[index] in [
"<pad>", "<sos>", "<unk>", "<eos>"
]:
continue
print(EN.vocab.itos[index], end=' ')
# print('test ...', word)
print()
print("[!] End the testing ...")
break
def main():
args = parse_arguments()
n_vocab = params.n_vocab
n_layer = params.n_layer
n_hidden = params.n_hidden
n_embed = params.n_embed
n_batch = args.n_batch
temperature = params.temperature
train_path = params.train_path
assert torch.cuda.is_available()
print("loading_data...")
# 训练时加载处理好的词典（如果有的话）
if os.path.exists("vocab.json"):
vocab = Vocabulary()
with open('vocab.json', 'r') as fp:
vocab.stoi = json.load(fp)
for key, value in vocab.stoi.items():
vocab.itos.append(key)
else:
vocab = build_vocab(train_path, n_vocab)
# save vocab
with open('vocab.json', 'w') as fp:
json.dump(vocab.stoi, fp)
train_X, train_y, train_K = load_data(train_path, vocab)
train_loader = get_data_loader(train_X, train_y, train_K, n_batch)
print("successfully loaded")
encoder = Encoder(n_vocab, n_embed, n_hidden, n_layer, vocab).cuda()
Kencoder = KnowledgeEncoder(n_vocab, n_embed, n_hidden, n_layer,
vocab).cuda()
manager = Manager(n_hidden, n_vocab, temperature).cuda()
decoder = Decoder(n_vocab, n_embed, n_hidden, n_layer, vocab).cuda()
if args.restore:
encoder = init_model(encoder, restore=params.encoder_restore)
Kencoder = init_model(Kencoder, restore=params.Kencoder_restore)
manager = init_model(manager, restore=params.manager_restore)
decoder = init_model(decoder, restore=params.decoder_restore)
# ToDo：目前的代码所有的embedding都是独立的，可以参考transformer源码使用直接赋值的方法共享参数：
#if emb_src_trg_weight_sharing:
# self.encoder.src_word_emb.weight = self.decoder.trg_word_emb.weight
model = [encoder, Kencoder, manager, decoder]
parameters = list(encoder.parameters()) + list(Kencoder.parameters()) + \
list(manager.parameters()) + list(decoder.parameters())
optimizer = optim.Adam(parameters, lr=args.lr)
# pre_train knowledge manager
print("start pre-training")
pre_train(model, optimizer, train_loader, args)
print("start training")
train(model, optimizer, train_loader, args)
# save final model
save_models(model, params.all_restore)
def main():
begin = time.time()
viser = Visualizer()
args = parse_arguments()
hidden_size = 256
embed_size = 128
assert torch.cuda.is_available()
print("[!] preparing dataset...")
train_iter, test_iter, val_iter, corpus = load_dataset(args.batch_size)
zh_size, en_size = corpus.zh_size, corpus.en_size
print("[!] Instantiating models...")
encoder = Encoder(en_size,
embed_size,
hidden_size,
n_layers=2,
dropout=0.5)
decoder = Decoder(embed_size,
hidden_size,
zh_size,
n_layers=1,
dropout=0.5)
seq2seq = Seq2Seq(encoder, decoder).cuda()
optimizer = optim.Adam(seq2seq.parameters(), lr=args.lr)
print("[!] Begin to train")
for e in range(1, args.epochs + 1):
train(e, seq2seq, optimizer, train_iter, zh_size, args.grad_clip,
corpus, viser, begin)
# val_loss = evaluate(seq2seq, val_iter, zh_size, corpus)
# print("[Epoch:%d] val_loss:%5.3f | val_pp:%5.2fS"
# % (e, val_loss, math.exp(val_loss)))
# Save the model if the validation loss is the best we've seen so far.
# if not best_val_loss or val_loss < best_val_loss:
if True:
print("[!] best model occur, saving model...")
# If find the new best model, try to plot text
if not os.path.isdir(".save"):
os.makedirs(".save")
torch.save(seq2seq.state_dict(), '.save/seq2seq_%d.pt' % (e))
# plot text from the test dataset
# only plot 1 sentences
test_batch = random.choice(train_iter)
src = test_batch.get_src()[:, 0].unsqueeze(0).cuda()
trg = test_batch.get_target()[:, 0].cuda()
pre = seq2seq.eval_forward(src, 100, corpus.zh_stoi['<sos>'])
src_sent = ' '.join([corpus.en_itos[int(word)] for word in src[0]])
trg_sent = ' '.join([corpus.zh_itos[int(word)] for word in trg])
pre_sent = ' '.join([corpus.zh_itos[int(word)] for word in pre])
viser.plot_text(e, src_sent, pre_sent, trg_sent)
# test_loss = evaluate(seq2seq, test_iter, zh_size, corpus)
# print("[TEST] Final loss:%5.2f" % test_loss)
print('[!] Training over')
def main(args):
torch.multiprocessing.set_start_method('spawn')
torch.distributed.init_process_group(backend="nccl")
with open(args.config_path, 'r') as file:
config = AttrDict(json.load(file))
set_seed(config.seed + torch.distributed.get_rank())
train_data_csv, test_data_csv = train_test_split(
config.train_data_csv_path, config.n_test_experiments)
train_image_ids, train_labels = get_data(train_data_csv, is_train=True)
train_transform = TrainTransform(config.crop_size)
train_dataset = CellsDataset(config.train_images_dir, train_image_ids,
train_labels, train_transform)
test_image_ids, test_labels = get_data(test_data_csv, is_train=True)
test_dataset = CellsDataset(config.train_images_dir, test_image_ids,
test_labels)
if torch.distributed.get_rank() == 0:
print(
f'Train size: {len(train_dataset)}, test_size: {len(test_dataset)}'
)
encoder = Encoder(config.n_image_channels, config.n_emedding_channels,
config.n_classes, config.encoder_model,
config.encoder_pretrained, config.encoder_dropout,
config.encoder_scale)
if config.restore_checkpoint_path is not None:
state_dict = torch.load(config.restore_checkpoint_path,
map_location='cpu')
encoder.load_state_dict(state_dict, strict=False)
decoder = Decoder(config.n_emedding_channels, config.n_image_channels,
config.n_classes, config.decoder_n_channels)
trainer = Trainer(encoder=encoder,
decoder=decoder,
optimizer_params={
'lr': config.lr,
'weight_decay': config.weight_decay,
'warmap': config.warmap,
'amsgrad': config.amsgrad
},
amp_params={
'opt_level': config.opt_level,
'loss_scale': config.loss_scale
},
rank=args.local_rank,
n_jobs=config.n_jobs)
trainer.train(train_data=train_dataset,
n_epochs=config.n_epochs,
batch_size=config.batch_size,
test_data=test_dataset,
best_checkpoint_path=config.best_checkpoint_path)
def build_model(vocab_size,
load_checkpoint=False,
checkpoint_epoch=-1,
print_module=True):
hidden_size = config['MODEL']['HIDDEN_SIZE']
attn_method = config['MODEL']['ATTN_METHOD']
num_encoder_layers = config['MODEL']['N_ENCODER_LAYERS']
dropout = config['MODEL']['DROPOUT']
encoder = Encoder(vocab_size,
hidden_size,
num_encoder_layers,
dropout=dropout)
decoder = Decoder(hidden_size,
vocab_size,
attn_method,
num_encoder_layers,
dropout=dropout)
model = Seq2Seq(encoder=encoder,
decoder=decoder,
max_length=config['LOADER']['MAX_LENGTH'],
tie_weights=config['MODEL']['TIE_WEIGHTS'])
if print_module:
print(model)
if load_checkpoint is True and os.path.exists(CHECKPOINT_PATH) is True:
# load checkpoint
prefix = config['TRAIN']['PREFIX']
model_path = None
if checkpoint_epoch >= 0:
model_path = '%s%s_%d' % (CHECKPOINT_PATH, prefix,
checkpoint_epoch)
else:
# use last checkpoint
checkpoints = []
for root, dirs, files in os.walk(CHECKPOINT_PATH):
for f_name in files:
f_name_sp = f_name.split('_')
if len(f_name_sp) == 2:
checkpoints.append(int(f_name_sp[1]))
if len(checkpoints) > 0:
model_path = '%s%s_%d' % (CHECKPOINT_PATH, prefix,
max(checkpoints))
if model_path is not None and os.path.exists(model_path):
if IMPORT_FROM_CUDA:
loaded = torch.load(model_path,
map_location=lambda storage, loc: storage)
else:
loaded = torch.load(model_path)
model.load_state_dict(loaded)
print('Load %s' % model_path)
# print('Seq2Seq parameters:')
# for name, param in model.state_dict().items():
# print(name, param.size())
if USE_CUDA:
model = model.cuda()
return model
def main(args):
# Is GPU usable?
assert torch.cuda.is_available()
# load encoder decoder vocab
logger.debug("Loading Vocabulary...")
# encoder vocaluraly
with open(args.en_vocab_path, "rb") as f:
en_vocab = pickle.load(f)
logger.debug("Encoder vocab size: {}".format(len(en_vocab)))
# decoder vocaburaly
with open(args.de_vocab_path, "rb") as f:
de_vocab = pickle.load(f)
logger.debug("Decoder vocab size: {}".format(len(de_vocab)))
en_size, de_size = len(en_vocab), len(de_vocab)
logger.debug("[source_vocab]:%d [target_vocab]:%d" % (en_size, de_size))
# setting train and val dataloader
logger.debug("Preparing dataset...")
train_iter = get_dataset(args.train_path, en_vocab, de_vocab,
args.batch_size, args.shuffle, args.num_workers)
val_iter = get_dataset(args.val_path, en_vocab, de_vocab, args.batch_size,
args.shuffle, args.num_workers)
# setting seq2seq model
logger.debug("Instantiating models...")
encoder = Encoder(en_size,
args.embed_dim,
args.hidden_dim,
n_layers=args.en_n_layers,
dropout=args.en_dropout)
decoder = Decoder(args.embed_dim,
args.hidden_dim,
de_size,
n_layers=args.de_n_layers,
dropout=args.de_dropout)
seq2seq = Seq2Seq(encoder, decoder).cuda()
if args.pre_trained_path is not None:
seq2seq.load_state_dict(torch.load(args.pre_trained_path))
logger.debug("Load pre trained model: {0}".format(
args.pre_trained_path))
optimizer = optim.Adam(seq2seq.parameters(), lr=args.lr)
logger.debug(seq2seq)
# Training and validation model
best_val_loss = None
for epoch in range(1, args.epochs + 1):
train(epoch, seq2seq, optimizer, train_iter, de_size, args.grad_clip,
en_vocab, de_vocab)
val_loss = evaluate(seq2seq, val_iter, de_size, en_vocab, de_vocab)
logger.debug("VAL LOSS: {0:.5f} (epoch={1})".format(val_loss, epoch))
# Save the model if the validation loss is the best we've seen so far.
if (best_val_loss is None) or (val_loss < best_val_loss):
logger.debug("save model (epoch={0})".format(epoch))
torch.save(seq2seq.state_dict(), args.save_model_path)
best_val_loss = val_loss
def main(config, tr_stream):
# Create Theano variables
logger.info('Creating theano variables')
source_char_seq = tensor.lmatrix('source_char_seq')
source_sample_matrix = tensor.btensor3('source_sample_matrix')
source_char_aux = tensor.bmatrix('source_char_aux')
source_word_mask = tensor.bmatrix('source_word_mask')
target_char_seq = tensor.lmatrix('target_char_seq')
target_char_aux = tensor.bmatrix('target_char_aux')
target_char_mask = tensor.bmatrix('target_char_mask')
target_sample_matrix = tensor.btensor3('target_sample_matrix')
target_word_mask = tensor.bmatrix('target_word_mask')
target_resample_matrix = tensor.btensor3('target_resample_matrix')
target_prev_char_seq = tensor.lmatrix('target_prev_char_seq')
target_prev_char_aux = tensor.bmatrix('target_prev_char_aux')
target_bos_idx = tr_stream.trg_bos
target_space_idx = tr_stream.space_idx['target']
src_vocab = pickle.load(open(config['src_vocab'], 'rb'))
logger.info('Building RNN encoder-decoder')
encoder = BidirectionalEncoder(config['src_vocab_size'], config['enc_embed'], config['src_dgru_nhids'],
config['enc_nhids'], config['src_dgru_depth'], config['bidir_encoder_depth'])
decoder = Decoder(config['trg_vocab_size'], config['dec_embed'], config['trg_dgru_nhids'], config['trg_igru_nhids'],
config['dec_nhids'], config['enc_nhids'] * 2, config['transition_depth'], config['trg_igru_depth'],
config['trg_dgru_depth'], target_space_idx, target_bos_idx)
representation = encoder.apply(source_char_seq, source_sample_matrix, source_char_aux,
source_word_mask)
cost = decoder.cost(representation, source_word_mask, target_char_seq, target_sample_matrix,
target_resample_matrix, target_char_aux, target_char_mask,
target_word_mask, target_prev_char_seq, target_prev_char_aux)
# Set up model
logger.info("Building model")
training_model = Model(cost)
# Set extensions
logger.info("Initializing extensions")
# Reload model if necessary
extensions = [LoadNMT(config['saveto'])]
# Initialize main loop
logger.info("Initializing main loop")
main_loop = MainLoop(
model=training_model,
algorithm=None,
data_stream=None,
extensions=extensions
)
for extension in main_loop.extensions:
extension.main_loop = main_loop
main_loop._run_extensions('before_training')
char_embedding = encoder.decimator.apply(source_char_seq.T, source_sample_matrix, source_char_aux.T)
embedding(Model(char_embedding), src_vocab)
