如何找到训练得到的模型？
lzhaozi opened this issue · 2 comments
我按照README的指示跑了一个基于combined的数据集的mutual learning的训练，最后得到了一个stdout***.txt，一个ckpt.pth和一个tensorboard文件夹，但是找不到.pth.tar模型文件。请问程序是没有保存模型文件么？还是说模型文件在ckpt.pth里，需要一些操作来提取出来？
谢谢。我是PyTorch初学者。。
命令是：
python script/experiment/train_ml.py -d '((2,), (3,))' -r 1 --num_models 2 --dataset combined --ids_per_batch 32 --ims_per_id 4 --normalize_feature false -gm 0.3 -glw 1 -llw 0 -idlw 0 -pmlw 0 -gdmlw 1 -ldmlw 0 --base_lr 2e-4 --lr_decay_type exp --exp_decay_at_epoch 151 --total_epochs 300
您好，多谢关注！
模型在ckpt.pth中，保存checkpoint的代码是save_ckpt(modules_optims, ep + 1, 0, cfg.ckpt_file)，参见这里，其中modules_optims = models + optimizers，参见这里。这里有两个模型在训练，还有两个optimizer，所以modules_optims是长度为4的list。根据save_ckpt的实现，存到磁盘的是一个dict，其中有一项是state_dicts，即modules_optims的参数，所以想要取得第一个模型的参数，可以这样做：
map_location = lambda storage, loc: storage
ckpt = torch.load(ckpt_file, map_location=map_location)
model_weight = ckpt['state_dicts'][0]