def main():
# 0. env preparation
place = paddle.fluid.CUDAPlace(0)
with fluid.dygraph.guard(place):
# 1. create model
if args.net == "pspnet":
model = PSPNet()
elif args.net == "unet":
model = UNet()
else:
raise NotImplementedError(
f"args.net: {args.net} is not Supported!")
# 2. load pretrained model
para_state_dict, _ = fluid.load_dygraph(args.model_folder + '/' +
args.net)
model.set_dict(para_state_dict)
model.eval()
# 3. read test image list
# 4. create transforms for test image, transform should be same as training
train_augmentation = Augmentation(image_size=256)
dataloader = BasicDataLoader(image_folder=args.image_folder,
image_list_file=args.image_list_file,
transform=train_augmentation,
shuffle=False)
# 5. loop over list of images
counter = 0
for im, _ in dataloader():
# 6. read image and do preprocessing
# 7. image to variable
counter += 1
im = im[np.newaxis, :, :, :]
im = to_variable(im)
# NHWC -> NCHW
im = fluid.layers.transpose(im, (0, 3, 1, 2))
pred = model(im)
pred = fluid.layers.softmax(pred, axis=1)
pred = fluid.layers.argmax(pred, axis=1)
pred = pred.numpy()
pred = np.squeeze(pred).astype('uint8')
# 8. call inference func
# if args.method == 'resize':
# pred = inference_resize(im, pred)
# elif args.method == 'sliding':
# pass
# else:
# raise Exception("Unexpected method '{}'".format(args.method))
# 9. save results
# NCHW -> NHWC
im = fluid.layers.transpose(im, (0, 2, 3, 1))
save_images(im.numpy(), pred, counter)
val_loader = DataLoader(TinySegData(phase='val'),
batch_size=1,
shuffle=False,
num_workers=0)
print("data loading finished...")
model = PSPNet(n_classes=6).to(device)
criterion = torch.nn.CrossEntropyLoss()
mkdirs = lambda x: os.makedirs(x, exist_ok=True)
model.load_state_dict(torch.load(
"./ckpt_seg/epoch_79_iou0.85.pth")) # the model storing path
# eval
model.eval()
for j, (images, seg_gts, rets) in enumerate(val_loader):
if j % 100 == 0:
print('{} sets finished...'.format(j))
# load data to device
images = images.to(device)
seg_gts = seg_gts.to(device)
# get prediction
seg_logit = model(images)
loss_seg = criterion(seg_logit, seg_gts.long())
seg_preds = torch.argmax(seg_logit, dim=1)
seg_preds_np = seg_preds.detach().cpu().numpy()
seg_gts_np = seg_gts.cpu().numpy()
def train(epo_num=50, show_vgg_params=False):
vis = visdom.Visdom()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg_model = VGGNet(requires_grad=True, show_params=show_vgg_params)
fcn_model = PSPNet(n_classes=2)
fcn_model = fcn_model.to(device)
criterion = nn.BCELoss().to(device)
optimizer = optim.SGD(fcn_model.parameters(), lr=1e-2, momentum=0.7)
all_train_iter_loss = []
all_test_iter_loss = []
# start timing
prev_time = datetime.now()
for epo in range(epo_num):
train_loss = 0
fcn_model.train()
for index, (img, lidar, label,
color_label) in enumerate(train_dataloader):
img = img.to(device)
lidar = lidar.to(device)
label = label.to(device)
optimizer.zero_grad()
output, out_cls = fcn_model(img)
output = torch.sigmoid(
output) # output.shape is torch.Size([4, 2, 160, 160])
loss = criterion(output, label)
loss.backward()
iter_loss = loss.item()
all_train_iter_loss.append(iter_loss)
train_loss += iter_loss
optimizer.step()
output_np = output.cpu().detach().numpy().copy()
output_np = np.argmin(output_np, axis=1)
bag_msk_np = label.cpu().detach().numpy().copy()
bag_msk_np = np.argmin(bag_msk_np, axis=1)
if np.mod(index, 15) == 0:
print('epoch {}, {}/{},train loss is {}'.format(
epo, index, len(train_dataloader), iter_loss))
vis.images(output_np[:, None, :, :],
win='train_pred',
opts=dict(title='train prediction'))
vis.images(bag_msk_np[:, None, :, :],
win='train_label',
opts=dict(title='label'))
vis.line(all_train_iter_loss,
win='train_iter_loss',
opts=dict(title='train iter loss'))
test_loss = 0
fcn_model.eval()
for index, (img, lidar, label,
color_label) in enumerate(test_dataloader):
img = img.to(device)
lidar = lidar.to(device)
label = label.to(device)
with torch.no_grad():
optimizer.zero_grad()
output, out_cls = fcn_model(img)
output = torch.sigmoid(output)
loss = criterion(output, label)
iter_loss = loss.item()
all_test_iter_loss.append(iter_loss)
test_loss += iter_loss
output_np = output.cpu().detach().numpy().copy()
output_np = np.argmin(output_np, axis=1)
bag_msk_np = label.cpu().detach().numpy().copy()
bag_msk_np = np.argmin(bag_msk_np, axis=1)
if np.mod(index, 15) == 0:
print(
r'Testing... Open http://localhost:8097/ to see test result.'
)
vis.images(output_np[:, None, :, :],
win='test_pred',
opts=dict(title='test prediction'))
vis.images(bag_msk_np[:, None, :, :],
win='test_label',
opts=dict(title='label'))
vis.line(all_test_iter_loss,
win='test_iter_loss',
opts=dict(title='test iter loss'))
cur_time = datetime.now()
h, remainder = divmod((cur_time - prev_time).seconds, 3600)
m, s = divmod(remainder, 60)
time_str = "Time %02d:%02d:%02d" % (h, m, s)
prev_time = cur_time
print('epoch train loss = %f, epoch test loss = %f, %s' %
(train_loss / len(train_dataloader),
test_loss / len(test_dataloader), time_str))
if np.mod(epo, 5) == 0:
torch.save(fcn_model, 'checkpoints/fcn_model_{}.pt'.format(epo))
print('saveing checkpoints/fcn_model_{}.pt'.format(epo))
