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