sampler=sampler,
num_workers=20)
nhead = 4
score_model = ScoreModel(kdd_dataset.unknown_token + 1,
1024,
1024,
use_bert=use_bert).cuda()
image_encoder = ImageEncoder(input_dim=2048, output_dim=1024, nhead=nhead)
image_encoder.load_pretrained_weights(
path='../user_data/image_encoder_large.pth')
image_encoder = image_encoder.cuda()
# text_generator = TextGenerator(score_model.embed.num_embeddings).cuda()
# score_model = ScoreModel(30522, 256, num_heads=1).cuda()
# category_embedding = CategoryEmbedding(256).cuda()
optimizer = Adam(score_model.get_params() + image_encoder.get_params())
if start_epoch > 0 and local_rank == 0:
checkpoints = torch.load(
os.path.join(checkpoints_dir,
'model-epoch{}.pth'.format(start_epoch)))
score_model.load_state_dict(checkpoints['score'])
image_encoder.load_state_dict(checkpoints['item'])
# text_generator.load_state_dict(checkpoints['generator'])
optimizer.load_state_dict(checkpoints['optimizer'])
print("load checkpoints")
# generator = iterate_minibatches(iters=(30 - start_epoch) * len(loader), batch_size=256, num_workers=8, root_dir='/home/dingyuhui/dataset/kdd-data', use_bert=use_bert)
scheduler = ExponentialLR(optimizer, 0.95, last_epoch=start_epoch - 1)
score_model = nn.parallel.DistributedDataParallel(
score_model,
num_workers=15)
nhead = 4
text_encoder = TextEncoder(kdd_dataset.unknown_token + 1,
1024,
256,
use_bert=use_bert).cuda()
image_encoder = ImageEncoder(input_dim=2048, output_dim=1024, nhead=nhead)
image_encoder.load_pretrained_weights(
path='../user_data/image_encoder_large.pth')
image_encoder = image_encoder.cuda()
score_model = ScoreModel(1024, 256).cuda()
# text_generator = TextGenerator(text_encoder.embed.num_embeddings).cuda()
# score_model = ScoreModel(30522, 256, num_heads=1).cuda()
# category_embedding = CategoryEmbedding(256).cuda()
optimizer = Adam(image_encoder.get_params() + text_encoder.get_params() +
score_model.get_params())
if start_epoch > 0 and local_rank == 0:
checkpoints = torch.load(
os.path.join(checkpoints_dir,
'model-epoch{}.pth'.format(start_epoch)), 'cpu')
text_encoder.load_state_dict(checkpoints['query'])
image_encoder.load_state_dict(checkpoints['item'])
score_model.load_state_dict(checkpoints['score'])
# text_generator.load_state_dict(checkpoints['generator'])
optimizer.load_state_dict(checkpoints['optimizer'])
print("load checkpoints")
# generator = iterate_minibatches(iters=(30 - start_epoch) * len(loader), batch_size=256, num_workers=8, root_dir='/home/dingyuhui/dataset/kdd-data', use_bert=use_bert)
scheduler = ExponentialLR(optimizer, 0.95, last_epoch=start_epoch - 1)
