def run(self):
encoder = utils.init_encoder(self.encoding)
traces = utils.gen_traces(self)
def _(loops):
for _ in range(loops):
for trace in traces:
encoder.put(trace)
encoder.encode()
yield _
model = args.model(src=None, trg=mscoco_dataset, args=args)
if args.mode == "train":
logger.info(str(model))
if args.load_encoder_from is not None:
if args.gpu > -1:
with torch.cuda.device(args.gpu):
encoder = torch.load(
str(args.model_path / args.load_encoder_from) + '.pt',
map_location=lambda storage, loc: storage.cuda())
else:
encoder = torch.load(str(args.model_path / args.load_encoder_from) +
'.pt',
map_location=lambda storage, loc: storage)
init_encoder(model, encoder)
logger.info("Pretrained encoder loaded.")
if args.load_from is not None:
if args.gpu > -1:
with torch.cuda.device(args.gpu):
model.load_state_dict(torch.load(
str(args.model_path / args.load_from) + '.pt',
map_location=lambda storage, loc: storage.cuda()),
strict=False) # load the pretrained models.
else:
model.load_state_dict(torch.load(
str(args.model_path / args.load_from) + '.pt',
map_location=lambda storage, loc: storage),
strict=False) # load the pretrained models.
logger.info("Pretrained model loaded.")