def __init__(self): self.device = 'cuda' self.epochs = 10 # Networks self.model = FCN32s(21) # Criterion self.criterion = nn.CrossEntropyLoss(ignore_index=255) # Optimizer self.lr = 0.1 self.optimizer = optim.Adam(self.model.parameters(), lr = self.lr) self.lr_scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=18, gamma=0.1) # Dataset self.train_set = PascalVOC(transforms=transforms.Compose([ RandomCrop(170), ToTensor() ])) self.valid_set = PascalVOC(names_file='./VOCdevkit/VOC2007/ImageSets/Segmentation/val.txt', transforms=transforms.Compose([ RandomCrop(170), ToTensor() ])) self.train_loader = DataLoader(self.train_set, batch_size=4, num_workers=2, shuffle=True) self.valid_loader = DataLoader(self.valid_set, batch_size=4, num_workers=2, shuffle=False)
def main(args):
cuda = torch.cuda.is_available()
# dataset
train_loader, test_loader = return_data(args)
# model
model = FCN16s(n_class=2)
start_epoch = 0
start_iteration = 0
if args.resume:
checkpoint = torch.load(args.name)
model.load_state_dict(checkpoint)
else:
fcn32s = FCN32s(n_class=2)
state_dict = torch.load(args.fcn32)
try:
fcn32s.load_state_dict(state_dict)
except RuntimeError:
fcn32s.load_state_dict(state_dict['model_state_dict'])
model.copy_params_from_fcn32s(fcn32s)
if cuda:
model = model.cuda()
# optimizer
if args.optimizer == 'SGD':
optim = torch.optim.SGD(get_parameters(model, bias=False),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'adam':
optim = torch.optim.Adam(get_parameters(model, bias=False),
lr=args.lr,
betas=(0.9, 0.999))
# train
trainer = Trainer(cuda=cuda,
model=model,
optimizer=optim,
train_loader=train_loader,
val_loader=test_loader,
epochs=200,
size_average=False,
name=args.name,
loss=args.loss)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
if args.mode == 'train':
trainer.train()
else:
trainer.draw()
def __init__(self, args):
super(Trainer, self).__init__()
self.epoch = args.epoch
self.batch_size = args.batch_size
self.data_dir = args.data_dir
self.save_dir = args.save_dir
self.result_dir = args.result_dir
self.log_dir = args.log_dir
self.gpu_mode = args.gpu_mode
self.verbose = args.verbose
if args.model == 'fcn16s':
self.model = FCN16s()
elif args.model == 'fcn32s':
self.model = FCN32s()
elif args.model == 'fcn8s':
self.model = FCN8s()
elif args.model == 'pspnet':
self.model = PSPnet()
else:
print("No this model type")
exit(-1)
if self.gpu_mode:
self.model = self.model.cuda()
self.parameter = self.model.parameters()
self.optimizer = optim.Adam(self.parameter, lr=args.learning_rate)
self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer,
gamma=0.5)
self.train_dataloader = get_data_loader(self.data_dir,
self.batch_size,
split='train')
self.test_dataloader = get_data_loader(self.data_dir, 1, split='val')
# experiment_id = args.model + time.strftime('%m%d%H%m')
# self.writer = SummaryWriter(log_dir=self.log_dir + '/tboard_' + experiment_id)
self.loss = FocalLoss(gamma=1.25)
if args.pretrain != '':
self._load_pretrain(args.pretrain)
#! /usr/bin/env python
import torch
import torch.nn as nn
import torch.utils.model_zoo as model_zoo
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
import math
#from ImageNetClassNames import classNames.type(dtype)
from PIL import Image
import pdb
import json
from util import *
from fcn32s import FCN32s
use_cuda = torch.cuda.is_available()
print("use_cuda: {}".format(use_cuda))
dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
vgg16 = models.vgg16(True)
fcn = FCN32s(3)
fcn.copy_params_from_vgg16(vgg16)
with open('fileNames.json', 'r') as f:
allNames = json.load(f)
name = allNames['train'][0]
import torch
import torch.nn as nn
from dataset import PascalVOC
from augmentation import RandomCrop, ToTensor
from torchvision import transforms
from torch.utils.data import DataLoader
from fcn32s import FCN32s
fcn = FCN32s(21)
train_dataset = PascalVOC(transforms=transforms.Compose(
[RandomCrop(200), ToTensor()]))
train_dataloader = DataLoader(train_dataset, batch_size=2)
criterion = nn.CrossEntropyLoss(ignore_index=255)
for batch in train_dataloader:
print(batch['image'].size())
print(batch['label'].size())
print(batch['image'].dtype)
print(batch['label'].dtype)
out = fcn(batch['image'])
loss = criterion(out, batch['label'])
print(loss)
break
def main(gpu=-1, batch_size=1, iterations=100000,
lr=1e-10, out='result', resume=''):
# prepare datasets
def transform(in_data):
img, label = in_data
vgg_subtract_bgr = np.array(
[103.939, 116.779, 123.68], np.float32)[:, None, None]
img -= vgg_subtract_bgr
img = pad(img, max_size=(512, 512), bg_value=0)
label = pad(label, max_size=(512, 512), bg_value=-1)
return img, label
train_data = VOCSemanticSegmentationDataset(mode='train')
test_data = VOCSemanticSegmentationDataset(mode='val')
extend(train_data, transform)
extend(test_data, transform)
# set up FCN32s
n_class = 21
model = FCN32s(n_class=n_class)
if gpu != -1:
model.to_gpu(gpu)
chainer.cuda.get_device(gpu).use()
# prepare an optimizer
optimizer = chainer.optimizers.MomentumSGD(lr=lr, momentum=0.99)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(rate=0.0005))
# prepare iterators
train_iter = chainer.iterators.SerialIterator(
train_data, batch_size=batch_size)
test_iter = chainer.iterators.SerialIterator(
test_data, batch_size=1, repeat=False, shuffle=False)
updater = training.StandardUpdater(train_iter, optimizer, device=gpu)
trainer = training.Trainer(updater, (iterations, 'iteration'), out=out)
val_interval = 3000, 'iteration'
log_interval = 100, 'iteration'
trainer.extend(
TestModeEvaluator(test_iter, model, device=gpu), trigger=val_interval)
# reporter related
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(extensions.PrintReport(
['iteration', 'main/time',
'main/loss', 'validation/main/loss',
'main/accuracy', 'validation/main/accuracy',
'main/accuracy_cls', 'validation/main/accuracy_cls',
'main/iu', 'validation/main/iu',
'main/fwavacc', 'validation/main/fwavacc']),
trigger=log_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
# visualize training
trainer.extend(
extensions.PlotReport(
['main/loss', 'validation/main/loss'],
trigger=log_interval, file_name='loss.png')
)
trainer.extend(
extensions.PlotReport(
['main/accuracy', 'validation/main/accuracy'],
trigger=log_interval, file_name='accuracy.png')
)
trainer.extend(
extensions.PlotReport(
['main/accuracy_cls', 'validation/main/accuracy_cls'],
trigger=log_interval, file_name='accuracy_cls.png')
)
trainer.extend(
extensions.PlotReport(
['main/iu', 'validation/main/iu'],
trigger=log_interval, file_name='iu.png')
)
trainer.extend(
extensions.PlotReport(
['main/fwavacc', 'validation/main/fwavacc'],
trigger=log_interval, file_name='fwavacc.png')
)
trainer.extend(
SemanticSegmentationVisReport(
range(10), # visualize outputs for the first 10 data of test_data
test_data,
model,
n_class=n_class,
predict_func=model.extract # use FCN32s.extract to get a score map
),
trigger=val_interval, invoke_before_training=True)
trainer.extend(extensions.dump_graph('main/loss'))
if resume:
chainer.serializers.load_npz(osp.expanduser(resume), trainer)
trainer.run()
def main():
parser = argparse.ArgumentParser(
description='ChainerCV Semantic Segmentation example with FCN')
parser.add_argument('--batch_size',
'-b',
type=int,
default=1,
help='Number of images in each mini-batch')
parser.add_argument('--iteration',
'-i',
type=int,
default=50000,
help='Number of iteration to carry out')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--lr',
'-l',
type=float,
default=1e-10,
help='Learning rate of the optimizer')
parser.add_argument('--out',
'-o',
default='result',
help='Directory to output the result')
parser.add_argument('--resume',
'-r',
default='',
help='Resume the training from snapshot')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batch_size))
print('# iteration: {}'.format(args.iteration))
print('')
batch_size = args.batch_size
iteration = args.iteration
gpu = args.gpu
lr = args.lr
out = args.out
resume = args.resume
# prepare datasets
def transform(in_data):
img, label = in_data
vgg_subtract_bgr = np.array([103.939, 116.779, 123.68],
np.float32)[:, None, None]
img -= vgg_subtract_bgr
img = transforms.pad(img, max_size=(512, 512), bg_value=0)
label = transforms.pad(label, max_size=(512, 512), bg_value=-1)
return img, label
train_data = VOCSemanticSegmentationDataset(mode='train')
test_data = VOCSemanticSegmentationDataset(mode='val')
train_data = TransformDataset(train_data, transform)
test_data = TransformDataset(test_data, transform)
# set up FCN32s
n_class = 21
model = FCN32s(n_class=n_class)
if gpu != -1:
model.to_gpu(gpu)
chainer.cuda.get_device(gpu).use()
# prepare an optimizer
optimizer = chainer.optimizers.MomentumSGD(lr=lr, momentum=0.99)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(rate=0.0005))
# prepare iterators
train_iter = chainer.iterators.SerialIterator(train_data,
batch_size=batch_size)
test_iter = chainer.iterators.SerialIterator(test_data,
batch_size=1,
repeat=False,
shuffle=False)
updater = training.StandardUpdater(train_iter, optimizer, device=gpu)
trainer = training.Trainer(updater, (iteration, 'iteration'), out=out)
val_interval = 3000, 'iteration'
log_interval = 100, 'iteration'
trainer.extend(TestModeEvaluator(test_iter, model, device=gpu),
trigger=val_interval)
# reporter related
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(extensions.PrintReport([
'iteration', 'main/time', 'main/loss', 'validation/main/loss',
'main/accuracy', 'validation/main/accuracy', 'main/accuracy_cls',
'validation/main/accuracy_cls', 'main/iu', 'validation/main/iu',
'main/fwavacc', 'validation/main/fwavacc'
]),
trigger=log_interval)
trainer.extend(extensions.ProgressBar(update_interval=10))
# visualize training
trainer.extend(
extensions.PlotReport(['main/loss', 'validation/main/loss'],
trigger=log_interval,
file_name='loss.png'))
trainer.extend(
extensions.PlotReport(['main/accuracy', 'validation/main/accuracy'],
trigger=log_interval,
file_name='accuracy.png'))
trainer.extend(
extensions.PlotReport(
['main/accuracy_cls', 'validation/main/accuracy_cls'],
trigger=log_interval,
file_name='accuracy_cls.png'))
trainer.extend(
extensions.PlotReport(['main/iu', 'validation/main/iu'],
trigger=log_interval,
file_name='iu.png'))
trainer.extend(
extensions.PlotReport(['main/fwavacc', 'validation/main/fwavacc'],
trigger=log_interval,
file_name='fwavacc.png'))
def vis_transform(in_data):
vgg_subtract_bgr = np.array([103.939, 116.779, 123.68],
np.float32)[:, None, None]
img, label = in_data
img += vgg_subtract_bgr
img, label = transforms.chw_to_pil_image_tuple((img, label),
indices=[0, 1])
return img, label
trainer.extend(
SemanticSegmentationVisReport(
range(10), # visualize outputs for the first 10 data of test_data
test_data,
model,
n_class=n_class,
predict_func=model.predict, # a function to predict output
vis_transform=vis_transform),
trigger=val_interval,
invoke_before_training=True)
trainer.extend(extensions.dump_graph('main/loss'))
if resume:
chainer.serializers.load_npz(osp.expanduser(resume), trainer)
trainer.run()