updates
This commit is contained in:
parent
2950f4c816
commit
8e327e3bd0
1
train.py
1
train.py
|
@ -204,6 +204,7 @@ def train():
|
|||
model.nc = nc # attach number of classes to model
|
||||
model.arc = opt.arc # attach yolo architecture
|
||||
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
|
||||
torch_utils.model_info(model, report='summary') # 'full' or 'summary'
|
||||
nb = len(dataloader)
|
||||
|
|
Loading…
Reference in New Issue