pytorch加载预训练模型并保存输出的特征文件为.npy

1. 构建模型和保存特征

model = build_model() # 构建模型.
model_weight = torch.load('weight_path.pth', map_location='cpu') # 加载预训练权重.
missing, unexpected = model.load_state_dict(model_weight , strict=True) # 将权重加载到模型.
print(f"Loading weights for model: missing: {len(missing)}, unexpected: {len(unexpected)}.")
model.eval()
with torch.no_grad():
	for iter, data in enumerate(train_dataloader):
		feature = model(input)
		feature = feature.cpu().numpy()
		numpy.save(f'feature/{feature_name}.npy', feature)

2. 加载特征文件

# 加载已经抽取好的特征.
feature = numpy.load('path.npy') # 格式是 ndarray
feature = torch.tensor(feature) # 转tensor
feature = feature.type(torch.float32) # 改变精度(可选)
feature = feature.to(device) # 加载到gpu

3. 保存模型

如果模型训练好了，保存模型的权重。

save_checkpoint_path = os.path.join(args.save_checkpoint_path, 'model-best.pth')
if args.gpus: # 多卡
    torch.save(model.module.state_dict(), save_checkpoint_path)
else: # 单卡
    torch.save(model.state_dict(), save_checkpoint_path)