예제 #1
0
def main(params):

    parser = argparse.ArgumentParser()
    parser.add_argument('--num_epochs',
                        type=int,
                        default=300,
                        help='Number of epochs to train for')
    parser.add_argument('--epoch_start_i',
                        type=int,
                        default=0,
                        help='Start counting epochs from this number')
    parser.add_argument('--checkpoint_step',
                        type=int,
                        default=5,
                        help='How often to save checkpoints (epochs)')
    parser.add_argument('--learning_rate',
                        type=float,
                        default=0.01,
                        help='learning rate used for train')
    parser.add_argument('--cuda',
                        type=str,
                        default='0',
                        help='GPU ids used for training')
    parser.add_argument('--save_model_path',
                        type=str,
                        default=None,
                        help='path to save model')
    parser.add_argument('--pretrained_model_path',
                        type=str,
                        default=None,
                        help='path to pretrained model')

    args = parser.parse_args(params)

    # create dataset and dataloader
    dataloader_train = DataLoader(train(input_transform, target_transform),
                                  num_workers=1,
                                  batch_size=2,
                                  shuffle=True)

    # build model
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
    model = DANet(nclass=2, backbone='resnet50', aux=False, se_loss=False)
    model = model.cuda()

    # build optimizer
    optimizer = torch.optim.RMSprop(model.parameters(), args.learning_rate)

    # load pretrained model if exists
    if args.pretrained_model_path is not None:
        print('load model from %s ...' % args.pretrained_model_path)
        model.module.load_state_dict(torch.load(args.pretrained_model_path))
        print('Done!')
    # train
    train(args, model, optimizer, dataloader_train)
예제 #2
0
파일: eval.py 프로젝트: zzz2010/Contrib
def get_model(args):
    model = DANet('DANet',
                  backbone=args.backbone,
                  num_classes=args.num_classes,
                  batch_size=1,
                  dilated=args.dilated,
                  multi_grid=args.multi_grid,
                  multi_dilation=args.multi_dilation)
    return model
예제 #3
0
def main(params):
    # basic parameters
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint_path',
                        type=str,
                        default=None,
                        required=True,
                        help='The path to the pretrained weights of model')
    parser.add_argument(
        '--crop_height',
        type=int,
        default=640,
        help='Height of cropped/resized input image to network')
    parser.add_argument('--crop_width',
                        type=int,
                        default=640,
                        help='Width of cropped/resized input image to network')
    parser.add_argument('--data',
                        type=str,
                        default='/path/to/data',
                        help='Path of training data')
    parser.add_argument('--batch_size',
                        type=int,
                        default=1,
                        help='Number of images in each batch')
    parser.add_argument('--context_path',
                        type=str,
                        default="resnet101",
                        help='The context path model you are using.')
    parser.add_argument('--cuda',
                        type=str,
                        default='0',
                        help='GPU ids used for training')
    parser.add_argument('--use_gpu',
                        type=bool,
                        default=True,
                        help='Whether to user gpu for training')
    parser.add_argument('--num_classes',
                        type=int,
                        default=2,
                        help='num of object classes (with void)')
    args = parser.parse_args(params)

    # create dataset and dataloader

    dataloader = DataLoader(test(input_transform, target_transform),
                            num_workers=1,
                            batch_size=1,
                            shuffle=True)
    # build model
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
    model = DANet(nclass=2, backbone='resnet50', aux=False, se_loss=False)
    model = model.cuda()

    # load pretrained model if exists
    print('load model from %s ...' % args.checkpoint_path)
    model.load_state_dict(torch.load(args.checkpoint_path))
    print('Done!')

    # test
    eval(model, dataloader, args)