def train(options):
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
num_processors = max(torch.cuda.device_count(), 1)
mt_model = SenSim(text_processor=text_processor, enc_layer=options.encoder_layer, embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
if options.pretrained_path is not None:
pret = Seq2Seq.load(Seq2Seq, options.pretrained_path, tok_dir=options.tokenizer_path)
mt_model.init_from_lm(pret)
print("Model initialization done!")
optimizer = build_optimizer(mt_model, options.learning_rate, warump_steps=options.warmup)
trainer = SenSimTrainer(model=mt_model, mask_prob=options.mask_prob, optimizer=optimizer, clip=options.clip,
fp16=options.fp16)
pin_memory = torch.cuda.is_available()
mt_train_loader = SenSimTrainer.get_mt_train_data(mt_model, num_processors, options, pin_memory)
src_neg_data = dataset.MassDataset(batch_pickle_dir=options.src_neg,
max_batch_capacity=num_processors * options.total_capacity * 5,
max_batch=num_processors * options.batch * 5,
pad_idx=mt_model.text_processor.pad_token_id(), keep_pad_idx=False,
max_seq_len=options.max_seq_len, keep_examples=False)
dst_neg_data = dataset.MassDataset(batch_pickle_dir=options.dst_neg,
max_batch_capacity=num_processors * options.total_capacity * 5,
max_batch=num_processors * options.batch * 5,
pad_idx=mt_model.text_processor.pad_token_id(), keep_pad_idx=False,
max_seq_len=options.max_seq_len, keep_examples=False)
src_neg_loader = data_utils.DataLoader(src_neg_data, batch_size=1, shuffle=True, pin_memory=pin_memory)
dst_neg_loader = data_utils.DataLoader(dst_neg_data, batch_size=1, shuffle=True, pin_memory=pin_memory)
mt_dev_loader = None
if options.mt_dev_path is not None:
mt_dev_loader = SenSimTrainer.get_mt_dev_data(mt_model, options, pin_memory, text_processor, trainer, )
step, train_epoch = 0, 1
trainer.best_loss = 1000000
while options.step > 0 and step < options.step:
print("train epoch", train_epoch)
step = trainer.train_epoch(mt_train_iter=mt_train_loader, max_step=options.step, mt_dev_iter=mt_dev_loader,
saving_path=options.model_path, step=step, src_neg_iter=src_neg_loader,
dst_neg_iter=dst_neg_loader)
train_epoch += 1
def train(options):
lex_dict = None
if options.dict_path is not None:
lex_dict = get_lex_dict(options.dict_path)
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
image_captioner = Seq2Seq.load(ImageCaptioning, options.pretrained_path, tok_dir=options.tokenizer_path)
txt2ImageModel = Caption2Image(text_processor=text_processor, enc_layer=options.encoder_layer,
embed_dim=options.embed_dim, intermediate_dim=options.intermediate_layer_dim)
print("Model initialization done!")
# We assume that the collator function returns a list with the size of number of gpus (in case of cpus,
collator = dataset.ImageTextCollator()
num_batches = max(1, torch.cuda.device_count())
optimizer = build_optimizer(txt2ImageModel, options.learning_rate, warump_steps=options.warmup)
trainer = Caption2ImageTrainer(model=txt2ImageModel, caption_model=image_captioner, mask_prob=options.mask_prob,
optimizer=optimizer,
clip=options.clip,
beam_width=options.beam_width, max_len_a=options.max_len_a,
max_len_b=options.max_len_b, len_penalty_ratio=options.len_penalty_ratio,
fp16=options.fp16, mm_mode=options.mm_mode)
pin_memory = torch.cuda.is_available()
img_train_loader = ImageMTTrainer.get_img_loader(collator, dataset.ImageCaptionDataset, options.train_path,
txt2ImageModel, num_batches, options, pin_memory,
lex_dict=lex_dict)
img_dev_loader = ImageMTTrainer.get_img_loader(collator, dataset.ImageCaptionDataset, options.dev_path,
txt2ImageModel, num_batches, options, pin_memory,
lex_dict=lex_dict,
shuffle=False, denom=2)
step, train_epoch = 0, 1
while options.step > 0 and step < options.step:
print("train epoch", train_epoch)
step = trainer.train_epoch(img_data_iter=img_train_loader, img_dev_data_iter=img_dev_loader,
max_step=options.step, lex_dict=lex_dict,
saving_path=options.model_path, step=step)
train_epoch += 1
def build_model(options):
model = Seq2Seq.load(ImageCaptioning,
options.model_path,
tok_dir=options.tokenizer_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
num_gpu = torch.cuda.device_count()
generator = BeamDecoder(model,
beam_width=options.beam_width,
max_len_a=options.max_len_a,
max_len_b=options.max_len_b,
len_penalty_ratio=options.len_penalty_ratio)
if options.fp16:
generator = amp.initialize(generator, opt_level="O2")
if num_gpu > 1:
generator = DataParallelModel(generator)
return generator, model.text_processor
def test_albert_seq2seq_init(self):
path_dir_name = os.path.dirname(os.path.realpath(__file__))
data_path = os.path.join(path_dir_name, "sample.txt")
with tempfile.TemporaryDirectory() as tmpdirname:
processor = TextProcessor()
processor.train_tokenizer([data_path],
vocab_size=1000,
to_save_dir=tmpdirname,
languages={
"<en>": 0,
"<fa>": 1
})
seq2seq = Seq2Seq(text_processor=processor)
src_inputs = torch.tensor([[
1, 2, 3, 4, 5,
processor.pad_token_id(),
processor.pad_token_id()
], [1, 2, 3, 4, 5, 6, processor.pad_token_id()]])
tgt_inputs = torch.tensor(
[[6, 8, 7,
processor.pad_token_id(),
processor.pad_token_id()],
[6, 8, 7, 8, processor.pad_token_id()]])
src_mask = (src_inputs != processor.pad_token_id())
tgt_mask = (tgt_inputs != processor.pad_token_id())
src_langs = torch.tensor([[0], [0]]).squeeze()
tgt_langs = torch.tensor([[1], [1]]).squeeze()
seq_output = seq2seq(src_inputs,
tgt_inputs,
src_mask,
tgt_mask,
src_langs,
tgt_langs,
log_softmax=True)
assert list(seq_output.size()) == [5, processor.vocab_size()]
seq_output = seq2seq(src_inputs, tgt_inputs, src_mask, tgt_mask,
src_langs, tgt_langs)
assert list(seq_output.size()) == [5, processor.vocab_size()]
[text_processor.lang_id(sentences[sid].strip().split(" ")[0])])
yield sid, source_tokenized, torch.LongTensor(
tids), candidates, src_lang, torch.LongTensor(target_langs)
if __name__ == "__main__":
parser = get_option_parser()
(options, args) = parser.parse_args()
print("Loading text processor...")
text_processor = TextProcessor(options.tokenizer_path)
num_processors = max(torch.cuda.device_count(), 1)
print("Loading model...")
model = Seq2Seq.load(Seq2Seq,
options.model,
tok_dir=options.tokenizer_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
num_gpu = torch.cuda.device_count()
assert num_gpu <= 1
if options.fp16:
model = amp.initialize(model, opt_level="O2")
max_capacity = options.total_capacity * 1000000
with torch.no_grad(), open(options.output, "w") as writer:
print("Loading data...")
with open(options.sens, "rb") as fp, open(options.data, "rb") as fp2:
sentences = marshal.load(fp)
src2dst_dict = marshal.load(fp2)
def train(options):
lex_dict = None
if options.dict_path is not None:
lex_dict = get_lex_dict(options.dict_path)
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
num_processors = max(torch.cuda.device_count(), 1)
if options.pretrained_path is not None:
print("Loading pretrained path", options.pretrained_path)
mt_model = Seq2Seq.load(ImageMassSeq2Seq,
options.pretrained_path,
tok_dir=options.tokenizer_path)
else:
mt_model = ImageMassSeq2Seq(
use_proposals=lex_dict is not None,
tie_embed=options.tie_embed,
text_processor=text_processor,
resnet_depth=options.resnet_depth,
lang_dec=options.lang_decoder,
enc_layer=options.encoder_layer,
dec_layer=options.decoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
if options.lm_path is not None:
lm = LM(text_processor=text_processor,
enc_layer=options.encoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
mt_model.init_from_lm(lm)
print("Model initialization done!")
# We assume that the collator function returns a list with the size of number of gpus (in case of cpus,
collator = dataset.ImageTextCollator()
num_batches = max(1, torch.cuda.device_count())
if options.continue_train:
with open(os.path.join(options.pretrained_path, "optim"),
"rb") as fp:
optimizer = pickle.load(fp)
else:
optimizer = build_optimizer(mt_model,
options.learning_rate,
warump_steps=options.warmup)
trainer = ImageMTTrainer(model=mt_model,
mask_prob=options.mask_prob,
optimizer=optimizer,
clip=options.clip,
beam_width=options.beam_width,
max_len_a=options.max_len_a,
max_len_b=options.max_len_b,
len_penalty_ratio=options.len_penalty_ratio,
fp16=options.fp16,
mm_mode=options.mm_mode)
pin_memory = torch.cuda.is_available()
mt_train_loader = None
if options.mt_train_path is not None:
mt_train_loader = ImageMTTrainer.get_mt_train_data(
mt_model,
num_processors,
options,
pin_memory,
lex_dict=lex_dict)
mt_dev_loader = None
if options.mt_dev_path is not None:
mt_dev_loader = ImageMTTrainer.get_mt_dev_data(mt_model,
options,
pin_memory,
text_processor,
trainer,
lex_dict=lex_dict)
step, train_epoch = 0, 1
while options.step > 0 and step < options.step and train_epoch <= 10:
print("train epoch", train_epoch, "step:", step)
step = trainer.train_epoch(mt_train_iter=mt_train_loader,
max_step=options.step,
lex_dict=lex_dict,
mt_dev_iter=mt_dev_loader,
saving_path=options.model_path,
step=step,
save_opt=False)
train_epoch += 1
def train(options):
lex_dict = None
if options.dict_path is not None:
lex_dict = get_lex_dict(options.dict_path)
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
if options.pretrained_path is not None:
mt_model = Seq2Seq.load(ImageCaptioning,
options.pretrained_path,
tok_dir=options.tokenizer_path)
else:
mt_model = ImageCaptioning(
use_proposals=lex_dict is not None,
tie_embed=options.tie_embed,
text_processor=text_processor,
resnet_depth=options.resnet_depth,
lang_dec=options.lang_decoder,
enc_layer=options.encoder_layer,
dec_layer=options.decoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
if options.lm_path is not None:
lm = LM(text_processor=text_processor,
enc_layer=options.encoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
mt_model.init_from_lm(lm)
print("Model initialization done!")
# We assume that the collator function returns a list with the size of number of gpus (in case of cpus,
collator = dataset.ImageTextCollator()
num_batches = max(1, torch.cuda.device_count())
if options.continue_train:
with open(os.path.join(options.pretrained_path, "optim"),
"rb") as fp:
optimizer = pickle.load(fp)
else:
optimizer = build_optimizer(mt_model,
options.learning_rate,
warump_steps=options.warmup)
trainer = ImageCaptionTrainer(
model=mt_model,
mask_prob=options.mask_prob,
optimizer=optimizer,
clip=options.clip,
beam_width=options.beam_width,
max_len_a=options.max_len_a,
max_len_b=options.max_len_b,
len_penalty_ratio=options.len_penalty_ratio,
fp16=options.fp16,
mm_mode=options.mm_mode)
pin_memory = torch.cuda.is_available()
img_train_loader = ImageMTTrainer.get_img_loader(
collator,
dataset.ImageCaptionDataset,
options.train_path,
mt_model,
num_batches,
options,
pin_memory,
lex_dict=lex_dict)
img_dev_loader = ImageMTTrainer.get_img_loader(
collator,
dataset.ImageCaptionDataset,
options.dev_path,
mt_model,
num_batches,
options,
pin_memory,
lex_dict=lex_dict,
shuffle=False,
denom=2)
trainer.reference = None
if img_dev_loader is not None:
trainer.reference = []
generator = (trainer.generator.module if hasattr(
trainer.generator, "module") else trainer.generator)
for data in img_dev_loader:
for batch in data:
captions = [b["captions"] for b in batch]
for caption in captions:
refs = get_outputs_until_eos(
text_processor.sep_token_id(),
caption,
remove_first_token=True)
ref = [
generator.seq2seq_model.text_processor.tokenizer.
decode(ref.numpy()) for ref in refs
]
trainer.reference += ref
step, train_epoch = 0, 1
while options.step > 0 and step < options.step:
print("train epoch", train_epoch)
step = trainer.train_epoch(img_data_iter=img_train_loader,
img_dev_data_iter=img_dev_loader,
max_step=options.step,
lex_dict=lex_dict,
saving_path=options.model_path,
step=step)
train_epoch += 1