yukun-hh/trash-division
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

添加模型断点继续加载训练功能

Browse source
This commit is contained in:
yukun-hh 2026-04-21 15:57:13 +08:00
parent 968e108857
commit f8bb340a70
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
Train.py
View file

@ -153,7 +153,7 @@ if __name__ == '__main__':
model = Net(num_classes=4) # 根据你的 Net 类调整
#断点继续训练
if os.path.exists('best_model.pth'):
model.load_state_dict(torch.load('best_model.pth'))
model.load_state_dict(torch.load('best_model.pth',map_location=torch.device('cpu')))
model = model.to(device)
# 打印模型信息