如何用ImageNet预训练模型分类自己的数据集?
如图,我自己设计了一个模型,用ImageNet进行预训练(类别为1000),训练好的参数被保存在pkl文件中。我想把这个参数用来分类另一个数据集(类别为10),当我把pkl文件载入后,报了如图所示的错误。请问应该如何修改呢?而且分类不同,又该怎么改呢?
以下为载入时的代码
model = Micro()
train_loader, val_loader = prepare_dataloader(train, trn_idx,val_idx, data_root=train_root)
model = torch.nn .DataParallel ( model ).cuda( )
model.load_state_dict({k.replace( ' module.' ,'' ) :v for k,v in torch. load( ' predicet/checkpoint_99_epoch.pkl ' ).items( )})//checkpoint_99_epoch.pkl 为预训练好的参数
mode1. load_state_dict(checkpoint)