[Python知识库]模型融合的模板代码

val_preds = []
tst_preds = []

# 每个 fold 的模型选择最好的 epoch 的结果，做 3 次 tta 融合得到最终结果。
for fold, epoch in enumerate(CFG['used_epochs']):    
    model.load_state_dict(torch.load('../input/{}-fold-{}-{}/{}_fold_{}_{}'.format('-'.join(CFG['model_arch'].split('_')), fold, epoch,CFG['model_arch'], fold, epoch)))

    with torch.no_grad():
        for _ in range(CFG['tta']):
            tst_preds += [CFG['weights'][fold]*inference_one_epoch(model, tst_loader, device)]

tst_preds = np.sum(tst_preds, axis=0) 

del model
torch.cuda.empty_cache()
test['label'] = np.argmax(tst_preds, axis=1)
test.to_csv('submission.csv', index=False)

CFG = {
    'model_arch': 'tf_efficientnet_b4_ns',
    'img_size': 512,
    'valid_bs': 32,
    'num_workers': 4,
    'device': 'cuda:0',
    'tta': 3,
    'used_epochs': [5,5,5,5,5],
    'weights': [1,1,1,1]
}