def get_model(name, n_classes):
model = fcn8s(n_classes)
#model = unet(n_classes)
#model = model(n_classes=n_classes)
vgg16 = models.vgg16(pretrained=True)
model.init_vgg16_params(vgg16)
return model
#print(data.max(), data.min())
#print(masks.max(), masks.min())
#show_sample3(data, masks)
#print(labels)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=conf['batch_size'])
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
# torch.cuda.set_device(6)
torch.cuda.set_device(conf['gpu'])
print(device)
dataset_sizes = {'test': len(test_dataset)}
from ptsemseg.models.fcn import fcn8s
model = fcn8s(n_classes=1)
vgg16 = models.vgg16(pretrained=True)
model = model.to(device)
if conf['num_gpus'] > 1:
model = torch.nn.DataParallel(model,
device_ids=range(torch.cuda.device_count()))
model_ft.load_state_dict(torch.load("./models/FCN8_ep4.net"))
from train import accuracy, iou
test(model, device, test_loader, epoch, dataset_sizes['test'], conf)
def train(model):
if model == 'unet':
model = unet(feature_scale=feature_scale,
n_classes=n_classes,
is_batchnorm=True,
in_channels=3,
is_deconv=True)
if model == 'segnet':
model = segnet(n_classes=n_classes, in_channels=3, is_unpooling=True)
if model == 'fcn32':
model = fcn32s(n_classes=n_classes)
vgg16 = models.vgg16(pretrained=True)
model.init_vgg16_params(vgg16)
if model == 'fcn16':
model = fcn16s(n_classes=n_classes)
vgg16 = models.vgg16(pretrained=True)
model.init_vgg16_params(vgg16)
if model == 'fcn8':
model = fcn8s(n_classes=n_classes)
vgg16 = models.vgg16(pretrained=True)
model.init_vgg16_params(vgg16)
pascal = pascalVOCLoader(data_path, is_transform=True, img_size=img_rows)
trainloader = data.DataLoader(pascal, batch_size=batch_size, num_workers=4)
if torch.cuda.is_available():
model.cuda(0)
optimizer = torch.optim.SGD(model.parameters(),
lr=l_rate,
momentum=0.99,
weight_decay=5e-4)
test_image, test_segmap = pascal[0]
test_image = Variable(test_image.unsqueeze(0).cuda(0))
vis = visdom.Visdom()
for epoch in range(n_epoch):
for i, (images, labels) in enumerate(trainloader):
if torch.cuda.is_available():
images = Variable(images.cuda(0))
labels = Variable(labels.cuda(0))
else:
images = Variable(images)
labels = Variable(labels)
optimizer.zero_grad()
outputs = model(images)
loss = cross_entropy2d(outputs, labels)
loss.backward()
optimizer.step()
if (i + 1) % 20 == 0:
print("Epoch [%d/%d] Loss: %.4f" %
(epoch + 1, n_epoch, loss.data[0]))
test_output = model(test_image)
predicted = pascal.decode_segmap(
test_output[0].cpu().data.numpy().argmax(0))
target = pascal.decode_segmap(test_segmap.numpy())
vis.image(test_image[0].cpu().data.numpy(),
opts=dict(title='Input' + str(epoch)))
vis.image(np.transpose(target, [2, 0, 1]),
opts=dict(title='GT' + str(epoch)))
vis.image(np.transpose(predicted, [2, 0, 1]),
opts=dict(title='Predicted' + str(epoch)))
torch.save(model, "unet_voc_" + str(feature_scale) + ".pkl")