changed enabled parameter
parent
77180d7087
commit
df838a42dd
3
train.py
3
train.py
|
@ -37,6 +37,7 @@ import torch.nn as nn
|
|||
import yaml
|
||||
from torch.optim import lr_scheduler
|
||||
from tqdm import tqdm
|
||||
from torch.cuda import amp
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
ROOT = FILE.parents[0] # YOLOv5 root directory
|
||||
|
@ -352,7 +353,7 @@ def train(hyp, opt, device, callbacks):
|
|||
maps = np.zeros(nc) # mAP per class
|
||||
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
|
||||
scheduler.last_epoch = start_epoch - 1 # do not move
|
||||
scaler = torch.amp.GradScaler(enabled=amp)
|
||||
scaler = amp.GradScaler(enabled=cuda)
|
||||
stopper, stop = EarlyStopping(patience=opt.patience), False
|
||||
compute_loss = ComputeLoss(model) # init loss class
|
||||
callbacks.run("on_train_start")
|
||||
|
|
Loading…
Reference in New Issue