Glenn Jocher
4 years ago
1 changed files with 2 additions and 2 deletions
+ 2
- 2
train.py
View File
| os.remove(f) | os.remove(f) | ||||
| # Create model | # Create model | ||||
| model = Model(opt.cfg, nc=data_dict['nc']).to(device) | |||||
| model = Model(opt.cfg, nc=nc).to(device) | |||||
| # Image sizes | # Image sizes | ||||
| gs = int(max(model.stride)) # grid size (max stride) | gs = int(max(model.stride)) # grid size (max stride) | ||||
| model.hyp = hyp # attach hyperparameters to model | model.hyp = hyp # attach hyperparameters to model | ||||
| model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou) | model.gr = 1.0 # giou loss ratio (obj_loss = 1.0 or giou) | ||||
| 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 | ||||
| model.names = data_dict['names'] | |||||
| model.names = names | |||||
| # Class frequency | # Class frequency | ||||
| labels = np.concatenate(dataset.labels, 0) | labels = np.concatenate(dataset.labels, 0) |
Loading…