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