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))