示例#1
0
if not os.path.exists('./output_img'):
    os.mkdir('./output_img')

parser, metadata = get_parser_with_args()
opt = parser.parse_args()

dev = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

test_loader = get_test_loaders(opt, batch_size=1)

path = 'weights/snunet-32.pt'  # the path of the model
model = torch.load(path)

model.eval()
index_img = 0
test_metrics = initialize_metrics()
with torch.no_grad():
    tbar = tqdm(test_loader)
    for batch_img1, batch_img2, labels in tbar:

        batch_img1 = batch_img1.float().to(dev)
        batch_img2 = batch_img2.float().to(dev)
        labels = labels.long().to(dev)

        cd_preds = model(batch_img1, batch_img2)

        cd_preds = cd_preds[-1]
        _, cd_preds = torch.max(cd_preds, 1)
        cd_preds = cd_preds.data.cpu().numpy()
        cd_preds = cd_preds.squeeze() * 255
示例#2
0
model = load_model(opt, dev)

criterion = get_criterion(opt)
# optimizer = optim.AdamW(model.parameters(), lr=opt.learning_rate)
optimizer = AdamW(model.parameters(), lr=opt.learning_rate)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.5)
"""
 Set starting values
"""
best_metrics = {'cd_f1scores': -1, 'cd_recalls': -1, 'cd_precisions': -1}
logging.info('STARTING training')
best_pre, best_recall, best_f1 = [-1, -1], [-1, -1], [-1, -1]
total_step = -1

for epoch in range(opt.epochs):
    train_metrics = initialize_metrics()
    val_metrics = initialize_metrics()
    """
    Begin Training
    """
    model.train()
    logging.info('SET model mode to train!')
    batch_iter = 0
    tbar = tqdm(train_loader)
    for batch_img1, batch_img2, labels in tbar:
        tbar.set_description("epoch {} info ".format(epoch) + str(batch_iter) +
                             " - " + str(batch_iter + opt.batch_size))
        batch_iter = batch_iter + opt.batch_size
        total_step += 1
        # Set variables for training
        batch_img1 = autograd.Variable(batch_img1).float().to(dev)
示例#3
0
from tqdm import tqdm

parser, metadata = get_parser_with_args()
opt = parser.parse_args()

dev = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

test_loader = get_test_loaders(opt)

path = 'tmp/checkpoint_epoch_64.pt'
model = torch.load(path)
criterion = get_criterion(opt)

model.eval()

val_metrics = initialize_metrics()
with torch.no_grad():
    tbar = tqdm(test_loader)
    for batch_img1, batch_img2, labels in tbar:
        # Set variables for training
        batch_img1 = autograd.Variable(batch_img1).float().to(dev)
        batch_img2 = autograd.Variable(batch_img2).float().to(dev)
        labels = autograd.Variable(labels).long().to(dev)

        # Get predictions and calculate loss
        cd_preds = model(batch_img1, batch_img2)

        cd_preds = cd_preds[-1]
        _, cd_preds = torch.max(cd_preds, 1)

        # Calculate and log other batch metrics