This commit is contained in:
Glenn Jocher 2019-11-20 13:33:25 -08:00
parent 2950f4c816
commit 8e327e3bd0
1 changed files with 2 additions and 1 deletions

View File

@ -204,6 +204,7 @@ def train():
model.nc = nc # attach number of classes to model model.nc = nc # attach number of classes to model
model.arc = opt.arc # attach yolo architecture model.arc = opt.arc # attach yolo architecture
model.hyp = hyp # attach hyperparameters to model model.hyp = hyp # attach hyperparameters to model
if hasattr(dataset, 'labels'):
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
torch_utils.model_info(model, report='summary') # 'full' or 'summary' torch_utils.model_info(model, report='summary') # 'full' or 'summary'
nb = len(dataloader) nb = len(dataloader)