def get_loaders(train_csv_path, test_csv_path):
from dataset import YoloDataset
IMAGE_SIZE = config.IMAGE_SIZE
train_dataset = YoloDataset(
train_csv_path,
transform=config.train_transforms,
grid_size=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
image_path=config.IMG_DIR,
label_path=config.LABEL_DIR,
anchors=config.ANCHORS,
)
test_dataset = YoloDataset(
test_csv_path,
transform=config.test_transforms,
grid_size=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
image_path=config.IMG_DIR,
label_path=config.LABEL_DIR,
anchors=config.ANCHORS,
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
pin_memory=config.PIN_MEMORY,
shuffle=True,
drop_last=False,
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
pin_memory=config.PIN_MEMORY,
shuffle=False,
drop_last=False,
)
train_eval_dataset = YoloDataset(
train_csv_path,
transform=config.test_transforms,
grid_size=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
image_path=config.IMG_DIR,
label_path=config.LABEL_DIR,
anchors=config.ANCHORS,
)
train_eval_loader = DataLoader(
dataset=train_eval_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
pin_memory=config.PIN_MEMORY,
shuffle=False,
drop_last=False,
)
return train_loader, test_loader, train_eval_loader
def eval(checkpoint, data_path, params):
# 数据
files, label, boxes = load_annotation(data_path, 'test')
eval_set = YoloDataset(paths=files,
bboxes=boxes,
labels=label,
params=params,
train=False)
eval_loader = DataLoader(eval_set,
batch_size=params.batch_size,
num_workers=params.num_gpus * 8,
shuffle=False)
# 模型
state_dict = torch.load(checkpoint)
model = Backbone()
model.load_state_dict(state_dict)
model = model.cuda()
# 损失
criterion = SumSquareError()
model.eval()
total_loss = 0
with torch.no_grad():
for iter, (img, annotation) in enumerate(eval_loader):
img = img.cuda()
annotation = annotation.cuda()
output = model(img)
loss = criterion(output, annotation).item()
total_loss += loss * len(img)
print(f'evaluate loss: {total_loss / len(eval_set)}')
def do_test(model,
images_path,
labels_path,
batch_size=32,
progress_callback=None):
size = 416
t = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
data = DataLoader(YoloDataset(images_path, labels_path, t, size),
batch_size=batch_size,
shuffle=True,
num_workers=mp.cpu_count())
count = [0] * 80
correct = [0] * 80
with torch.no_grad():
for i, (local_batch, local_labels) in enumerate(data):
local_batch = local_batch.to(device)
outputs = model(local_batch)
detections = yolo.YoloV3.get_detections(outputs)
for j, detections in enumerate(utils.parse_detections(detections)):
detections = utils.non_max_suppression(
detections, confidence_threshold=0.2)
for target in utils.parse_labels(local_labels[0][j], size):
count[target['coco_idx']] += 1
for det in [
det for det in detections
if det.coco_idx == target['coco_idx']
]:
if utils.iou(target['bb'], det.bb) >= 0.5:
correct[target['coco_idx']] += 1
break
psum = 0
for j in range(80):
if count[j] == 0:
psum = -80
break
p = correct[j] / count[j]
psum += p
if progress_callback:
progress_callback(model,
batch_number=(i + 1),
batch_count=len(data),
map_score=psum / 80)
return psum / 80
def main():
model = YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE,
dtype=torch.float)
optimizer = optim.Adam(model.parameters(),
lr=config.LEARNING_RATE,
weight_decay=config.WEIGHT_DECAY)
loss_fn = YoloLoss().to(config.DEVICE)
scaler = torch.cuda.amp.GradScaler()
dataset = YoloDataset(table=table,
anchors=anchors,
transform=config.transform)
train_loader = DataLoader(dataset, batch_size=2, shuffle=True)
for epoch in (range(20)):
train_fn(train_loader, model, optimizer, loss_fn, scaler, anchors)
def main(options):
image_size = 416
t = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
data = DataLoader(YoloDataset(options.images, options.labels, t,
image_size),
batch_size=options.batch_size,
shuffle=True,
num_workers=mp.cpu_count())
model = yolo.load_model(train_dn=False,
train_yolo=True,
load_weights=False)
model.to(device)
# values from:
# [yolov3.cfg](https://github.com/pjreddie/darknet/blob/master/cfg/yolov3.cfg)
# optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9,
# dampening=0, weight_decay=0.0005)
optimizer = torch.optim.Adam(model.parameters())
# for param in model.parameters():
# print(param.requires_grad)
# # param.detach_()
print('num learnable params:', len(list(model.parameters())))
for epoch in range(options.epochs):
for i, (local_batch, local_labels) in enumerate(data):
# use the same input over and over again for debugging
for i in range(100):
optimizer.zero_grad()
local_batch = local_batch.to(device)
outputs = model(local_batch)
# detections = yolo.YoloV3.get_detections(outputs)
# print(detections.shape)
loss = yolo.YoloV3.get_loss(outputs, local_labels)
def test() -> None:
dataset = YoloDataset(table, anchors, shuffle=True)
train_loader = DataLoader(dataset, batch_size=1, shuffle=False)
myloss = YoloLoss()
itr = train_loader.__iter__()
image, targets = next(itr)
b_, c_, h_, w_ = image.shape
fakes = targets.copy()
for idx, f in enumerate(fakes):
b, c, h, w, m = f.shape
index = f.reshape(-1, 6)[..., 5:].type(torch.int64)
y = torch.zeros((b * c * h * w, 3)).scatter_(-1, index, 1)
y = y.reshape((b, c, h, w, -1))
pred = torch.cat([f[..., :5], y], dim=-1)
pred[..., 3:4] /= w_
pred[..., 4:5] /= h_
loss = myloss.forward(pred, targets[idx], anchors[idx])
def train(args):
yolo_dataset = YoloDataset(DATASET,
"faces.csv",
input_size=208,
transform=transforms.Compose([ToTensor()]))
dataset_indices = set([i for i in range(len(yolo_dataset))])
test_indices = random.sample(dataset_indices,
int(args.test_size * len(yolo_dataset)))
train_indices = list(dataset_indices - set(test_indices))
trainset = Subset(yolo_dataset, train_indices)
testset = Subset(yolo_dataset, train_indices)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
validloader = torch.utils.data.DataLoader(testset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
dataloaders = {TRAIN: trainloader, VALID: validloader}
model = MyResnet()
if CUDA:
model.cuda()
optimizer = optim.Adam(model.conv_addition.parameters(), lr=args.lr)
saver = CheckpointSaver(args.save_dir_name, max_checkpoints=3)
scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
# writer = SummaryWriter(os.path.join(args.save_dir_name, "log"))
train_losses, test_losses = [], []
for epoch in range(args.epochs):
for phase in dataloaders:
if phase == TRAIN:
scheduler.step()
model.train()
else:
model.eval()
epoch_avg = AVG()
logger.info(f"-----------------{phase.upper()}-----------------")
for i, data in enumerate(dataloaders[phase]):
img, y = data['img'], data['y']
if CUDA:
img = img.cuda()
y = y.cuda()
optimizer.zero_grad()
with torch.set_grad_enabled(phase == TRAIN):
pred = model(img)
y_ = pred.permute((0, 2, 3, 1))
y = y.permute((0, 2, 3, 1))
loss = loss_function(y_, y)
epoch_avg.add(loss.item())
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
logger.info(f"Epoch: {epoch}, batch: {i}, loss {loss.item()}")
logger.info(f"Epoch: {epoch}, average loss: {epoch_avg}")
if phase == TRAIN:
train_losses.append(str(epoch_avg))
else:
test_losses.append(str(epoch_avg))
# writer.add_scalar(f'{phase}_data/average_loss', str(epoch_avg), epoch)
if epoch % 20 == 0:
saver.save(model, optimizer, epoch)
with open(os.path.join(args.save_dir_name, "losses.txt"), 'w') as file:
file.write(str(train_losses))
file.write(str(test_losses))
def train(params, _run=None):
params = Params(params)
set_random_seeds(params.seed)
time_now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
params.save_root = params.save_root + f'/{params.project_name}_{time_now}_{params.version}'
os.makedirs(params.save_root, exist_ok=True)
logging.basicConfig(filename=f'{params.save_root}/{params.project_name}_{time_now}_{params.version}.log',
filemode='a', format='%{asctime}s - %(levalname)s: %(message)s')
if params.num_gpus == 0:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
logging.info(f'Available GPUs: {torch.cuda.device_count()}')
train2007, train_label_2007, train_bb_2007 = load_annotation(os.path.join(params.data_root, 'VOC2007'), 'trainval')
test2007, test_label_2007, test_bb_2007 = load_annotation(os.path.join(params.data_root, 'VOC2007'), 'test')
train2012, train_label_2012, train_bb_2012 = load_annotation(os.path.join(params.data_root, 'VOC2012'), 'trainval')
test2012, test_label_2012, test_bb_2012 = load_annotation(os.path.join(params.data_root, 'VOC2012'), 'test')
train_data = train2007+test2007+train2012
train_label = train_label_2007+test_label_2007+train_label_2012
train_bb = train_bb_2007 + test_bb_2007 + train_bb_2012
test_data = test2012
test_label = test_label_2012
test_bb = test_bb_2012
train_dataset = YoloDataset(train_data, train_bb, train_label, params, train=True)
eval_dataset = YoloDataset(test_data, test_bb, test_label, params, train=False)
train_loader = DataLoader(dataset=train_dataset, num_workers=params.num_gpus*8, batch_size=params.batch_size,
shuffle=True, drop_last=True, pin_memory=True)
eval_loader = DataLoader(dataset=eval_dataset, num_workers=1, batch_size=1,
shuffle=False, pin_memory=True)
model = Backbone()
last_step = 0
last_epoch = 0
if params.num_gpus > 0:
model = model.cuda()
if params.num_gpus > 1:
model = nn.DataParallel(model)
if params.optim == 'Adam':
optimizer = torch.optim.Adam(model.parameters(), lr=params.learning_rate)
else:
optimizer = torch.optim.SGD(model.parameters(), lr=params.learning_rate, momentum=0.9, nesterov=True, weight_decay=0.0005)
criterion = SumSquareError()
schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer, factor=0.5, verbose=True, patience=10)
epoch = 0
begin_epoch = max(0, last_epoch)
step = max(0, last_step)
best_loss = 1e6
logging.info('Begin to train...')
model.train()
import cv2 as cv
try:
for epoch in range(begin_epoch, params.epoch):
for iter, (img, annotation) in enumerate(train_loader):
output = model(img.cuda())
loss = criterion(output, annotation.cuda())
optimizer.zero_grad()
loss.backward()
optimizer.step()
if iter % params.save_interval == 0:
logging.info(f'{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} '
f'Train Epoch: {epoch} iter: {iter} loss: {loss.item()}')
step += 1
if epoch % params.eval_interval == 0:
model.eval()
epoch_loss = 0
with torch.no_grad():
for iter, (img, annotation) in enumerate(eval_loader):
output = model(img.cuda())
loss = criterion(output, annotation.cuda()).item()
epoch_loss += loss * len(img)
loss = epoch_loss / len(eval_dataset)
logging.info(f'{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} '
f'Eval Epoch: {epoch} loss: {loss}')
schedule.step(loss)
if loss < best_loss:
best_loss = loss
save_checkpoint(model, f'{params.save_root}/{epoch}_{step}.pth')
model.train()
except KeyboardInterrupt:
save_checkpoint(model, f'{params.save_root}/Interrupt_{epoch}_{step}.pth')
def valid_NB(datafile, cfgfile, weightfile, epoch, prefix='result'):
model = Darknet(cfgfile)
options = read_data_file(datafile)
data_root = options['valid']
class_root = options['names']
names = read_class_names(class_root)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
with open(data_root, 'r') as f:
lines = f.readlines()
valid_files = [item.strip() for item in lines]
model.load_weights(weightfile)
model = model.to(device)
model.eval()
NB_model = joblib.load('NB models/epoch_{}.pkl'.format(int(epoch)))
data = YoloDataset(data_root,
shape=(model.width, model.height),
transform=transforms.Compose([transforms.ToTensor()]),
train=False)
batch_size = 2
kwargs = {'num_workers': 4, 'pin_memory': True}
data_loader = DataLoader(data,
batch_size=batch_size,
shuffle=False,
**kwargs)
fs = [None] * model.num_classes
if not os.path.exists(prefix):
os.makedirs(prefix)
for i in range(model.num_classes):
filename = prefix + '/' + str(names[i]) + '.txt'
fs[i] = open(filename, 'w')
net_shape = (model.width, model.height)
conf_thresh = 0.005
nms_thresh = 0.45
fileIndex = 0
for index, (imgs, labels, org_w, org_h) in enumerate(data_loader):
imgs = imgs.to(device)
output = model(imgs)
batch_boxes = get_all_boxes(output,
net_shape,
conf_thresh,
model.num_classes,
device,
validation=True)
for i in range(len(batch_boxes)):
fileId = os.path.basename(valid_files[fileIndex]).split('.')[
0] # gei naive image name without suffix
w, h = float(org_w[i]), float(org_h[i])
# print(valid_files[fileIndex], '{}/{}'.format(fileIndex+1, len(data_loader) * batch_size))
fileIndex += 1
boxes = batch_boxes[i]
correct_yolo_boxes(boxes, w, h, model.width, model.height)
boxes = auto_thresh_nms(boxes, NB_model)
for box in boxes:
x1 = (box[0] - box[2] / 2.0) * w
y1 = (box[1] - box[3] / 2.0) * h
x2 = (box[0] + box[2] / 2.0) * w
y2 = (box[1] + box[3] / 2.0) * h
# 包含物体的概率，乘以每一类的概率
det_conf = box[4]
for j in range((len(box) - 5) // 2):
cls_conf = box[5 + 2 * j]
cls_id = int(box[5 + 2 * j + 1])
prob = det_conf * cls_conf
fs[cls_id].write('{:s} {:f} {:f} {:f} {:f} {:f}\n'.format(
fileId, prob, x1, y1, x2, y2))
for i in range(len(fs)):
fs[i].close()
def yolo_predictions(images):
batch_size = 1
confidence = 0.5
nms_thresh = 0.4
img_size = 256
classes = load_classes('coco.names')
CUDA = torch.cuda.is_available()
# Set up the neural network
print("Loading network.....")
model = Darknet('./yolov3.cfg')
model.load_weights('yolov3.weights')
print("Network successfully loaded")
model.net_info["height"] = 256
inp_dim = int(model.net_info["height"])
assert inp_dim % 32 == 0
assert inp_dim > 32
# If there's a GPU available, put the model on GPU
if CUDA:
_gpu = 0
torch.cuda.set_device(_gpu)
model.cuda()
print("Using GPU: {}".format(_gpu))
# Set the model in evaluation mode
model.eval()
# create dataset from images in the current folder
image_transform = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.ToTensor(),
transforms.Normalize((0., 0., 0.), (1, 1, 1))])
dataset = YoloDataset(images, transform=image_transform)
assert dataset
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
drop_last=False, shuffle=False, num_workers=4)
num_batches = len(dataloader)
dataloader = iter(dataloader)
output_dict = {}
# get YOLO predictions for images in current folder
for idx in tqdm(range(num_batches)):
data = dataloader.next()
imgs, filenames = data
# print(filenames)
if CUDA:
imgs = imgs.cuda()
with torch.no_grad():
predictions = model(imgs, CUDA)
predictions = non_max_suppression(predictions, confidence, nms_thresh)
for img, preds in zip(filenames, predictions):
img_preds_name = []
img_preds_id = []
img_bboxs = []
if preds is not None and len(preds) > 0:
for pred in preds:
pred_id = int(pred[-1])
pred_name = classes[pred_id]
bbox_x = pred[0] / img_size
bbox_y = pred[1] / img_size
bbox_width = (pred[2] - pred[0]) / img_size
bbox_height = (pred[3] - pred[1]) / img_size
img_preds_id.append(pred_id)
img_preds_name.append(pred_name)
img_bboxs.append([bbox_x.cpu().numpy(), bbox_y.cpu().numpy(),
bbox_width.cpu().numpy(), bbox_height.cpu().numpy()])
output_dict[img.split("/")[-1]] = [img_preds_name, img_preds_id, img_bboxs]
return output_dict