def test(path):
model = Model()
model.to("cuda:0")
model.eval()
checkpoint = torch.load("./model.pth")
model.load_state_dict(checkpoint["model"])
img = np.array(Image.open(path).resize([448, 448]))[np.newaxis]
img = np.transpose(img, axes=[0, 3, 1, 2]) / 255
img = torch.tensor(img, dtype=torch.float32).to("cuda:0")
preds = model(img).cpu().detach().numpy()
cell_h, cell_w = IMG_H / S, IMG_W / S
x, y = np.meshgrid(range(S), range(S))
preds_xywhs = []
for i in range(B):
preds_x = (preds[0, :, :, i * 4] + x) * cell_w
preds_y = (preds[0, :, :, i * 4 + 1] + y) * cell_h
preds_w = preds[0, :, :, i * 4 + 2] * IMG_W
preds_h = preds[0, :, :, i * 4 + 3] * IMG_H
preds_xywh = np.dstack((preds_x, preds_y, preds_w, preds_h))
preds_xywhs.append(preds_xywh)
preds_xywhs = np.dstack(preds_xywhs)
preds_xywhs = np.reshape(preds_xywhs, [-1, 4])
preds_class = preds[0, :, :, 10:]
preds_class = np.reshape(preds_class, [-1, 20])
preds_c = preds[0, :, :, 8:10]
preds_c = np.reshape(preds_c, [-1, 1])
max_arg = np.argmax(preds_c, axis=0)
print("max confidence: %f" % (preds_c[max_arg]))
max_arg_ = np.argmax(preds_class[int(max_arg // 2)])
print("class confidence: %f" % (preds_class[max_arg // 2, max_arg_]))
print("class category: %s" % (CLASSES[int(max_arg_)]))
Image.fromarray(
np.uint8(
draw_bboxes(np.array(Image.open(path).resize([448, 448])),
preds_xywhs[max_arg[0]:max_arg[0] + 1]))).show()
def train():
model = Model()
model.to("cuda:0")
Opt = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)
# checkpoint = torch.load("./model.pth")
# model.load_state_dict(checkpoint["model"])
# Opt.load_state_dict(checkpoint["Opt"])
for i in range(10000):
Opt.zero_grad()
imgs, targets = read_batch()
imgs = torch.tensor(imgs, dtype=torch.float32).to("cuda:0")
targets = torch.tensor(targets, dtype=torch.float32).to("cuda:0")
preds = model(imgs)
loss = make_loss(preds, targets)
loss.backward()
Opt.step()
if i % 10 == 0:
print("Iteration: %d, Loss: %f"%(i, loss))
if i % 10 == 0:
state = {'model':model.state_dict(), 'Opt':Opt.state_dict(), 'itr':i}
torch.save(state, "./model.pth")
train_indices, val_indices = indices[split:], indices[:split]
# Creating data samplers and loaders:
train_sampler = SubsetRandomSampler(train_indices)
valid_sampler = SubsetRandomSampler(val_indices)
train_loader = DataLoader(my_dataset,
batch_size=batch_size,
sampler=train_sampler)
validation_loader = DataLoader(my_dataset,
batch_size=batch_size,
sampler=valid_sampler)
# model, loss, and optimizer settings
model = Model()
model.to(device=device)
# print(model)
optimizer = torch.optim.Adam(model.parameters(),
lr=0.001,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=0.001,
amsgrad=False)
# training
# print(model.state_dict())
model = model.float()
loss_train = []
loss_val = []
latent = np.empty((1, 16, 16, 128))
