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