mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

pretrained model loading bug fix (#450)

Browse Source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2020-07-19 22:12:55 -07:00
parent 07a82f4d44
commit b569ed6d6b
1 changed files with 14 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
train.py
View File

@ -1,13 +1,12 @@
import argparse
import torch
import torch.distributed as dist
import torch.nn.functional as F
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torch.utils.data
from torch.utils.tensorboard import SummaryWriter
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.tensorboard import SummaryWriter
import test # import test.py to get mAP after each epoch
from models.yolo import Model
@ -131,7 +130,8 @@ def train(hyp, tb_writer, opt, device):
# load model
try:
ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items() if k in model.state_dict()}
ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items()
if k in model.state_dict() and model.state_dict()[k].shape == v.shape}
model.load_state_dict(ckpt['model'], strict=False)
except KeyError as e:
s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \
@ -187,7 +187,8 @@ def train(hyp, tb_writer, opt, device):
# Trainloader
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True,
cache=opt.cache_images, rect=opt.rect, local_rank=local_rank, world_size=opt.world_size)
cache=opt.cache_images, rect=opt.rect, local_rank=local_rank,
world_size=opt.world_size)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
@ -242,7 +243,8 @@ def train(hyp, tb_writer, opt, device):
if local_rank in [-1, 0]:
w = model.class_weights.cpu().numpy() * (1 - maps) ** 2 # class weights
image_weights = labels_to_image_weights(dataset.labels, nc=nc, class_weights=w)
dataset.indices = random.choices(range(dataset.n), weights=image_weights, k=dataset.n) # rand weighted idx
dataset.indices = random.choices(range(dataset.n), weights=image_weights,
k=dataset.n) # rand weighted idx
# Broadcast.
if local_rank != -1:
indices = torch.zeros([dataset.n], dtype=torch.int)
@ -431,7 +433,8 @@ if __name__ == '__main__':
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
parser.add_argument("--sync-bn", action="store_true", help="Use sync-bn, only avaible in DDP mode.")
# Parameter For DDP.
parser.add_argument('--local_rank', type=int, default=-1, help="Extra parameter for DDP implementation. Don't use it manually.")
parser.add_argument('--local_rank', type=int, default=-1,
help="Extra parameter for DDP implementation. Don't use it manually.")
opt = parser.parse_args()
last = get_latest_run() if opt.resume == 'get_last' else opt.resume # resume from most recent run