def main():
logger.auto_set_dir()
global args
parser = argparse.ArgumentParser()
parser.add_argument('--dataroot',
default='/home/hutao/lab/pytorchgo/example/ROAD/data',
help='Path to source dataset')
parser.add_argument('--batchSize',
type=int,
default=1,
help='input batch size')
parser.add_argument('--max_epoch',
type=int,
default=max_epoch,
help='Number of training iterations')
parser.add_argument('--optimizer',
type=str,
default='Adam',
help='Optimizer to use | SGD, Adam')
parser.add_argument('--lr',
type=float,
default=base_lr,
help='learning rate')
parser.add_argument('--momentum',
type=float,
default=0.99,
help='Momentum for SGD')
parser.add_argument('--beta1',
type=float,
default=0.9,
help='beta1 for adam. default=0.5')
parser.add_argument('--weight_decay',
type=float,
default=0.0005,
help='Weight decay')
parser.add_argument('--model', type=str, default='vgg16')
parser.add_argument('--gpu', type=int, default=1)
args = parser.parse_args()
print(args)
gpu = args.gpu
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
logger.info("random seed 1337")
torch.cuda.manual_seed(1337)
# Defining data loaders
kwargs = {
'num_workers': 4,
'pin_memory': True,
'drop_last': True
} if cuda else {}
train_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA(
'SYNTHIA',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True,
**kwargs)
val_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes(
'cityscapes',
args.dataroot,
split='val',
transform=True,
image_size=image_size),
batch_size=1,
shuffle=False)
target_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes(
'cityscapes',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True)
if cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if args.model == "vgg16":
model = origin_model = torchfcn.models.Seg_model(n_class=class_num)
vgg16 = torchfcn.models.VGG16(pretrained=True)
model.copy_params_from_vgg16(vgg16)
model_fix = torchfcn.models.Seg_model(n_class=class_num)
model_fix.copy_params_from_vgg16(vgg16)
for param in model_fix.parameters():
param.requires_grad = False
elif args.model == "deeplabv2": # TODO may have problem!
model = origin_model = torchfcn.models.Res_Deeplab(
num_classes=class_num, image_size=image_size)
saved_state_dict = model_zoo.load_url(Deeplabv2_restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
# print i_parts
if not class_num == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
model_fix = torchfcn.models.Res_Deeplab(num_classes=class_num,
image_size=image_size)
model_fix.load_state_dict(new_params)
else:
raise ValueError("only support vgg16, deeplabv2!")
netD = torchfcn.models.Domain_classifer(reverse=True)
netD.apply(weights_init)
model_summary([model, netD])
if cuda:
model = model.cuda()
netD = netD.cuda()
# Defining optimizer
if args.optimizer == 'SGD':
raise ValueError("SGD is not prepared well..")
optim = torch.optim.SGD([
{
'params': get_parameters(model, bias=False)
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2,
'weight_decay': args.weight_decay
},
],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'Adam':
if args.model == "vgg16":
optim = torch.optim.Adam([
{
'params': get_parameters(model, bias=False),
'weight_decay': args.weight_decay
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2,
'weight_decay': args.weight_decay
},
],
lr=args.lr,
betas=(args.beta1, 0.999))
elif args.model == "deeplabv2":
optim = torch.optim.Adam(origin_model.optim_parameters(args.lr),
lr=args.lr,
betas=(args.beta1, 0.999),
weight_decay=args.weight_decay)
else:
raise
else:
raise ValueError('Invalid optmizer argument. Has to be SGD or Adam')
optimD = torch.optim.Adam(netD.parameters(),
lr=dis_lr,
weight_decay=args.weight_decay,
betas=(0.7, 0.999))
optimizer_summary([optim, optimD])
trainer = MyTrainer_ROAD(cuda=cuda,
model=model,
model_fix=model_fix,
netD=netD,
optimizer=optim,
optimizerD=optimD,
train_loader=train_loader,
target_loader=target_loader,
val_loader=val_loader,
batch_size=args.batchSize,
image_size=image_size,
loss_print_interval=LOSS_PRINT_INTERVAL)
trainer.epoch = 0
trainer.iteration = 0
trainer.train()
def _init_():
if not os.path.exists('checkpoints'):
os.makedirs('checkpoints')
if not os.path.exists('checkpoints/' + args.exp_name):
os.makedirs('checkpoints/' + args.exp_name)
if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
os.system('cp main.py checkpoints' + '/' + args.exp_name + '/' +
'main.py.backup')
os.system('cp models/models.py checkpoints' + '/' + args.exp_name + '/' +
'models.py.backup')
_init_()
logger.auto_set_dir()
logger.info(str(args))
args.cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
logger.info('Using GPU : ' + str(torch.cuda.current_device()) + ' from ' +
str(torch.cuda.device_count()) + ' devices')
torch.cuda.manual_seed(args.seed)
else:
logger.info('Using CPU')
def train_batch(model, data):
[enc_nn, metric_nn, softmax_module] = model
def train(args):
logger.auto_set_dir()
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
# Setup Augmentations
data_aug = Compose([RandomRotate(10), RandomHorizontallyFlip()])
# Setup Dataloader
data_loader = get_loader(args.dataset)
data_path = get_data_path(args.dataset)
t_loader = data_loader(data_path,
is_transform=True,
img_size=(args.img_rows, args.img_cols),
epoch_scale=4,
augmentations=data_aug,
img_norm=args.img_norm)
v_loader = data_loader(data_path,
is_transform=True,
split='val',
img_size=(args.img_rows, args.img_cols),
img_norm=args.img_norm)
n_classes = t_loader.n_classes
trainloader = data.DataLoader(t_loader,
batch_size=args.batch_size,
num_workers=8,
shuffle=True)
valloader = data.DataLoader(v_loader,
batch_size=args.batch_size,
num_workers=8)
# Setup Metrics
running_metrics = runningScore(n_classes)
# Setup Model
from model_zoo.deeplabv1 import VGG16_LargeFoV
model = VGG16_LargeFoV(class_num=n_classes,
image_size=[args.img_cols, args.img_rows],
pretrained=True)
#model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
model.cuda()
# Check if model has custom optimizer / loss
if hasattr(model, 'optimizer'):
logger.warn("don't have customzed optimizer, use default setting!")
optimizer = model.module.optimizer
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=args.l_rate,
momentum=0.99,
weight_decay=5e-4)
optimizer_summary(optimizer)
if args.resume is not None:
if os.path.isfile(args.resume):
logger.info(
"Loading model and optimizer from checkpoint '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
logger.info("Loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
logger.info("No checkpoint found at '{}'".format(args.resume))
best_iou = -100.0
for epoch in tqdm(range(args.n_epoch), total=args.n_epoch):
model.train()
for i, (images, labels) in tqdm(enumerate(trainloader),
total=len(trainloader),
desc="training epoch {}/{}".format(
epoch, args.n_epoch)):
cur_iter = i + epoch * len(trainloader)
cur_lr = adjust_learning_rate(optimizer,
args.l_rate,
cur_iter,
args.n_epoch * len(trainloader),
power=0.9)
#if i > 10:break
images = Variable(images.cuda())
labels = Variable(labels.cuda())
optimizer.zero_grad()
outputs = model(images)
#print(np.unique(outputs.data[0].cpu().numpy()))
loss = CrossEntropyLoss2d_Seg(input=outputs,
target=labels,
class_num=n_classes)
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
logger.info("Epoch [%d/%d] Loss: %.4f, lr: %.7f" %
(epoch + 1, args.n_epoch, loss.data[0], cur_lr))
model.eval()
for i_val, (images_val, labels_val) in tqdm(enumerate(valloader),
total=len(valloader),
desc="validation"):
images_val = Variable(images_val.cuda(), volatile=True)
labels_val = Variable(labels_val.cuda(), volatile=True)
outputs = model(images_val)
pred = outputs.data.max(1)[1].cpu().numpy()
gt = labels_val.data.cpu().numpy()
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
logger.info("{}: {}".format(k, v))
running_metrics.reset()
if score['Mean IoU : \t'] >= best_iou:
best_iou = score['Mean IoU : \t']
state = {
'epoch': epoch + 1,
'mIoU': best_iou,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(),
}
torch.save(state,
os.path.join(logger.get_logger_dir(), "best_model.pkl"))
def main():
logger.auto_set_dir()
parser = argparse.ArgumentParser()
parser.add_argument('--dataroot',
default='/home/hutao/lab/pytorchgo/example/ROAD/data',
help='Path to source dataset')
parser.add_argument('--batchSize',
type=int,
default=1,
help='input batch size')
parser.add_argument('--num_iters',
type=int,
default=100000,
help='Number of training iterations')
parser.add_argument('--optimizer',
type=str,
default='Adam',
help='Optimizer to use | SGD, Adam')
parser.add_argument('--lr',
type=float,
default=1.0e-5,
help='learning rate')
parser.add_argument('--momentum',
type=float,
default=0.99,
help='Momentum for SGD')
parser.add_argument('--beta1',
type=float,
default=0.9,
help='beta1 for adam. default=0.5')
parser.add_argument('--weight_decay',
type=float,
default=0.0005,
help='Weight decay')
parser.add_argument(
'--interval_validate',
type=int,
default=3000,
help=
'Period for validation. Model is validated every interval_validate iterations'
)
parser.add_argument('--gpu', type=int, default=2)
args = parser.parse_args()
print(args)
gpu = args.gpu
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
logger.info("random seed 1337")
torch.cuda.manual_seed(1337)
# Defining data loaders
kwargs = {
'num_workers': 4,
'pin_memory': True,
'drop_last': True
} if cuda else {}
train_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA(
'SYNTHIA',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True,
**kwargs)
val_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes(
'cityscapes',
args.dataroot,
split='val',
transform=True,
image_size=image_size),
batch_size=1,
shuffle=False)
target_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes(
'cityscapes',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True)
# Defining models
start_epoch = 0
start_iteration = 0
if cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
model = torchfcn.models.Seg_model(n_class=class_num)
model_fix = torchfcn.models.Seg_model(n_class=class_num)
for param in model_fix.parameters():
param.requires_grad = False
netD = torchfcn.models.Domain_classifer(reverse=True)
netD.apply(weights_init)
model_summary(model_fix)
model_summary(netD)
vgg16 = torchfcn.models.VGG16(pretrained=True)
model.copy_params_from_vgg16(vgg16)
if cuda:
model = model.cuda()
netD = netD.cuda()
# Defining optimizer
if args.optimizer == 'SGD':
optim = torch.optim.SGD([
{
'params': get_parameters(model, bias=False)
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2,
'weight_decay': args.weight_decay
},
],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'Adam':
optim = torch.optim.Adam([
{
'params': get_parameters(model, bias=False)
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2
},
],
lr=args.lr,
betas=(args.beta1, 0.999))
else:
raise ValueError('Invalid optmizer argument. Has to be SGD or Adam')
optimD = torch.optim.Adam(netD.parameters(), lr=0.0001, betas=(0.7, 0.999))
trainer = MyTrainer_ROAD(cuda=cuda,
model=model,
model_fix=model_fix,
netD=netD,
optimizer=optim,
optimizerD=optimD,
train_loader=train_loader,
target_loader=target_loader,
val_loader=val_loader,
max_iter=args.num_iters,
batch_size=args.batchSize,
interval_validate=args.interval_validate,
image_size=image_size)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
def train(args):
logger.auto_set_dir()
from pytorchgo.utils.pytorch_utils import set_gpu
set_gpu(args.gpu)
# Setup Dataloader
from pytorchgo.augmentation.segmentation import SubtractMeans, PIL2NP, RGB2BGR, PIL_Scale, Value255to0, ToLabel
from torchvision.transforms import Compose, Normalize, ToTensor
img_transform = Compose([ # notice the order!!!
PIL_Scale(train_img_shape, Image.BILINEAR),
PIL2NP(),
RGB2BGR(),
SubtractMeans(),
ToTensor(),
])
label_transform = Compose([
PIL_Scale(train_img_shape, Image.NEAREST),
PIL2NP(),
Value255to0(),
ToLabel()
])
val_img_transform = Compose([
PIL_Scale(train_img_shape, Image.BILINEAR),
PIL2NP(),
RGB2BGR(),
SubtractMeans(),
ToTensor(),
])
val_label_transform = Compose([
PIL_Scale(train_img_shape, Image.NEAREST),
PIL2NP(),
ToLabel(),
# notice here, training, validation size difference, this is very tricky.
])
from pytorchgo.dataloader.pascal_voc_loader import pascalVOCLoader as common_voc_loader
train_loader = common_voc_loader(split="train_aug",
epoch_scale=1,
img_transform=img_transform,
label_transform=label_transform)
validation_loader = common_voc_loader(split='val',
img_transform=val_img_transform,
label_transform=val_label_transform)
n_classes = train_loader.n_classes
trainloader = data.DataLoader(train_loader,
batch_size=args.batch_size,
num_workers=8,
shuffle=True)
valloader = data.DataLoader(validation_loader,
batch_size=args.batch_size,
num_workers=8)
# Setup Metrics
running_metrics = runningScore(n_classes)
# Setup Model
from pytorchgo.model.deeplabv1 import VGG16_LargeFoV
from pytorchgo.model.deeplab_resnet import Res_Deeplab
model = Res_Deeplab(NoLabels=n_classes, pretrained=True, output_all=False)
from pytorchgo.utils.pytorch_utils import model_summary, optimizer_summary
model_summary(model)
def get_validation_miou(model):
model.eval()
for i_val, (images_val, labels_val) in tqdm(enumerate(valloader),
total=len(valloader),
desc="validation"):
if i_val > 5 and is_debug == 1: break
if i_val > 200 and is_debug == 2: break
#img_large = torch.Tensor(np.zeros((1, 3, 513, 513)))
#img_large[:, :, :images_val.shape[2], :images_val.shape[3]] = images_val
output = model(Variable(images_val, volatile=True).cuda())
output = output
pred = output.data.max(1)[1].cpu().numpy()
#pred = output[:, :images_val.shape[2], :images_val.shape[3]]
gt = labels_val.numpy()
running_metrics.update(gt, pred)
score, class_iou = running_metrics.get_scores()
for k, v in score.items():
logger.info("{}: {}".format(k, v))
running_metrics.reset()
return score['Mean IoU : \t']
model.cuda()
# Check if model has custom optimizer / loss
if hasattr(model, 'optimizer'):
logger.warn("don't have customzed optimizer, use default setting!")
optimizer = model.module.optimizer
else:
optimizer = torch.optim.SGD(model.optimizer_params(args.l_rate),
lr=args.l_rate,
momentum=0.99,
weight_decay=5e-4)
optimizer_summary(optimizer)
if args.resume is not None:
if os.path.isfile(args.resume):
logger.info(
"Loading model and optimizer from checkpoint '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
logger.info("Loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
logger.info("No checkpoint found at '{}'".format(args.resume))
best_iou = 0
logger.info('start!!')
for epoch in tqdm(range(args.n_epoch), total=args.n_epoch):
model.train()
for i, (images, labels) in tqdm(enumerate(trainloader),
total=len(trainloader),
desc="training epoch {}/{}".format(
epoch, args.n_epoch)):
if i > 10 and is_debug == 1: break
if i > 200 and is_debug == 2: break
cur_iter = i + epoch * len(trainloader)
cur_lr = adjust_learning_rate(optimizer,
args.l_rate,
cur_iter,
args.n_epoch * len(trainloader),
power=0.9)
images = Variable(images.cuda())
labels = Variable(labels.cuda())
optimizer.zero_grad()
outputs = model(images) # use fusion score
loss = CrossEntropyLoss2d_Seg(input=outputs,
target=labels,
class_num=n_classes)
#for i in range(len(outputs) - 1):
#for i in range(1):
# loss = loss + CrossEntropyLoss2d_Seg(input=outputs[i], target=labels, class_num=n_classes)
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
logger.info(
"Epoch [%d/%d] Loss: %.4f, lr: %.7f, best mIoU: %.7f" %
(epoch + 1, args.n_epoch, loss.data[0], cur_lr, best_iou))
cur_miou = get_validation_miou(model)
if cur_miou >= best_iou:
best_iou = cur_miou
state = {
'epoch': epoch + 1,
'mIoU': best_iou,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(),
}
torch.save(state,
os.path.join(logger.get_logger_dir(), "best_model.pth"))
args.mixup_alpha = 0.4
else:
args.mixup_alpha = 1.5
else:
if not args.manimixup:
args.mixup_alpha = 1.0
else:
args.mixup_alpha = 2.0
if (args.pointmixup and args.rot) and (not args.manimixup):
args.align = True
else:
args.align = False
# logger.auto_set_dir('d', "{}_{}".format(args.data, args.model))
logger.auto_set_dir('d', "{}".format(args.savename))
if args.data == 'modelnet40':
num_class = 40
dataset_cls = ModelNet40Cls
else:
raise NotImplementedError
args.epochs = int(args.epochs)
from tensorboardX import SummaryWriter
writer = SummaryWriter(comment=args.savename)
writer.add_text('args', str(args), 0)
transforms_test = d_utils.PointcloudToTensor()
if args.rot:
def main():
logger.auto_set_dir()
global args, best_prec1
import argparse
parser = argparse.ArgumentParser(description="PyTorch implementation of Temporal Segment Networks")
parser.add_argument('--dataset', type=str,default="something", choices=['something', 'jester', 'moments'])
parser.add_argument('--modality', type=str, default="RGB", choices=['RGB', 'Flow'])
parser.add_argument('--train_list', type=str, default="")
parser.add_argument('--val_list', type=str, default="")
parser.add_argument('--root_path', type=str, default="")
parser.add_argument('--store_name', type=str, default="")
# ========================= Model Configs ==========================
parser.add_argument('--arch', type=str, default="BNInception")
parser.add_argument('--num_segments', type=int, default=3)
parser.add_argument('--consensus_type', type=str, default='avg')
parser.add_argument('--k', type=int, default=3)
parser.add_argument('--dropout', '--do', default=0.8, type=float,
metavar='DO', help='dropout ratio (default: 0.5)')
parser.add_argument('--loss_type', type=str, default="nll",
choices=['nll'])
parser.add_argument('--img_feature_dim', default=256, type=int, help="the feature dimension for each frame")
# ========================= Learning Configs ==========================
parser.add_argument('--epochs', default=120, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('-b', '--batch_size', default=128, type=int,
metavar='N', help='mini-batch size (default: 256)')
parser.add_argument('--lr', '--learning-rate', default=0.001, type=float,
metavar='LR', help='initial learning rate')
parser.add_argument('--lr_steps', default=[50, 100], type=float, nargs="+",
metavar='LRSteps', help='epochs to decay learning rate by 10')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
help='momentum')
parser.add_argument('--weight-decay', '--wd', default=5e-4, type=float,
metavar='W', help='weight decay (default: 5e-4)')
parser.add_argument('--clip-gradient', '--gd', default=20, type=float,
metavar='W', help='gradient norm clipping (default: disabled)')
parser.add_argument('--no_partialbn', '--npb', default=False, action="store_true")
# ========================= Monitor Configs ==========================
parser.add_argument('--print-freq', '-p', default=20, type=int,
metavar='N', help='print frequency (default: 10)')
parser.add_argument('--eval-freq', '-ef', default=5, type=int,
metavar='N', help='evaluation frequency (default: 5)')
# ========================= Runtime Configs ==========================
parser.add_argument('-j', '--workers', default=30, type=int, metavar='N',
help='number of data loading workers (default: 4)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
help='evaluate model on validation set')
parser.add_argument('--snapshot_pref', type=str, default="")
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
help='manual epoch number (useful on restarts)')
parser.add_argument('--gpu', type=str, default='4')
parser.add_argument('--flow_prefix', default="", type=str)
parser.add_argument('--root_log', type=str, default='log')
parser.add_argument('--root_model', type=str, default='model')
parser.add_argument('--root_output', type=str, default='output')
args = parser.parse_args()
args.consensus_type = "TRN"
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
device_ids = [int(id) for id in args.gpu.split(',')]
assert len(device_ids) >1, "TRN must run with GPU_num > 1"
args.root_log = logger.get_logger_dir()
args.root_model = logger.get_logger_dir()
args.root_output = logger.get_logger_dir()
categories, args.train_list, args.val_list, args.root_path, prefix = datasets_video.return_dataset(args.dataset, args.modality)
num_class = len(categories)
args.store_name = '_'.join(['TRN', args.dataset, args.modality, args.arch, args.consensus_type, 'segment%d'% args.num_segments])
print('storing name: ' + args.store_name)
model = TSN(num_class, args.num_segments, args.modality,
base_model=args.arch,
consensus_type=args.consensus_type,
dropout=args.dropout,
img_feature_dim=args.img_feature_dim,
partial_bn=not args.no_partialbn)
crop_size = model.crop_size
scale_size = model.scale_size
input_mean = model.input_mean
input_std = model.input_std
policies = model.get_optim_policies()
train_augmentation = model.get_augmentation()
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)#TODO, , device_ids=[int(id) for id in args.gpu.split(',')]
if torch.cuda.is_available():
model.cuda()
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print(("=> loaded checkpoint '{}' (epoch {})"
.format(args.evaluate, checkpoint['epoch'])))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
cudnn.benchmark = True
# Data loading code
if args.modality != 'RGBDiff':
normalize = GroupNormalize(input_mean, input_std)
else:
normalize = IdentityTransform()
if args.modality == 'RGB':
data_length = 1
elif args.modality in ['Flow', 'RGBDiff']:
data_length = 5
train_loader = torch.utils.data.DataLoader(
TSNDataSet(args.root_path, args.train_list, num_segments=args.num_segments,
new_length=data_length,
modality=args.modality,
image_tmpl=prefix,
transform=torchvision.transforms.Compose([
train_augmentation,
Stack(roll=(args.arch in ['BNInception','InceptionV3'])),
ToTorchFormatTensor(div=(args.arch not in ['BNInception','InceptionV3'])),
normalize,
])),
batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
TSNDataSet(args.root_path, args.val_list, num_segments=args.num_segments,
new_length=data_length,
modality=args.modality,
image_tmpl=prefix,
random_shift=False,
transform=torchvision.transforms.Compose([
GroupScale(int(scale_size)),
GroupCenterCrop(crop_size),
Stack(roll=(args.arch in ['BNInception','InceptionV3'])),
ToTorchFormatTensor(div=(args.arch not in ['BNInception','InceptionV3'])),
normalize,
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
# define loss function (criterion) and optimizer
if args.loss_type == 'nll':
criterion = torch.nn.CrossEntropyLoss().cuda()
else:
raise ValueError("Unknown loss type")
for group in policies:
logger.info('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format(
group['name'], len(group['params']), group['lr_mult'], group['decay_mult']))
optimizer = torch.optim.SGD(policies,
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.evaluate:
validate(val_loader, model, criterion, 0)
return
log_training = open(os.path.join(args.root_log, '%s.csv' % args.store_name), 'w')
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch, args.lr_steps)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, log_training)
# evaluate on validation set
if (epoch + 1) % args.eval_freq == 0 or epoch == args.epochs - 1:
prec1 = validate(val_loader, model, criterion, (epoch + 1) * len(train_loader), log_training)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, is_best)
def main():
logger.auto_set_dir()
parser = argparse.ArgumentParser()
parser.add_argument(
'--dataroot',
default='/home/hutao/lab/pytorchgo/example/LSD-seg/data',
help='Path to source dataset')
parser.add_argument('--batchSize',
type=int,
default=1,
help='input batch size')
parser.add_argument('--num_iters',
type=int,
default=100000,
help='Number of training iterations')
parser.add_argument('--optimizer',
type=str,
default='Adam',
help='Optimizer to use | SGD, Adam')
parser.add_argument('--lr',
type=float,
default=1.0e-5,
help='learning rate')
parser.add_argument('--momentum',
type=float,
default=0.99,
help='Momentum for SGD')
parser.add_argument('--beta1',
type=float,
default=0.9,
help='beta1 for adam. default=0.5')
parser.add_argument('--weight_decay',
type=float,
default=0.0005,
help='Weight decay')
parser.add_argument(
'--interval_validate',
type=int,
default=500,
help=
'Period for validation. Model is validated every interval_validate iterations'
)
parser.add_argument(
'--resume',
default='',
help=
"path to the current checkpoint for resuming training. Do not specify if model has to be trained from scratch"
)
parser.add_argument('--method',
default='LSD',
help="Method to use for training | LSD, sourceonly")
parser.add_argument('--l1_weight', type=float, default=1, help='L1 weight')
parser.add_argument('--adv_weight',
type=float,
default=0.1,
help='Adv_weight')
parser.add_argument('--c_weight', type=float, default=0.1, help='C_weight')
parser.add_argument('--gpu', type=int, default=0)
args = parser.parse_args()
print(args)
gpu = args.gpu
out = logger.get_logger_dir()
resume = args.resume
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
# Defining data loaders
image_size = [640, 320]
kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}
train_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA(
'SYNTHIA',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True,
**kwargs)
val_loader = torch.utils.data.DataLoader(torchfcn.datasets.SYNTHIA(
'SYNTHIA',
args.dataroot,
split='val',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=False,
**kwargs)
target_loader = torch.utils.data.DataLoader(torchfcn.datasets.CityScapes(
'cityscapes',
args.dataroot,
split='train',
transform=True,
image_size=image_size),
batch_size=args.batchSize,
shuffle=True)
# Defining models
start_epoch = 0
start_iteration = 0
if args.method == 'sourceonly':
model = torchfcn.models.FCN8s_sourceonly(n_class=class_num)
elif args.method == 'LSD':
model = torchfcn.models.FCN8s_LSD(n_class=class_num)
netG = torchfcn.models._netG()
netD = torchfcn.models._netD()
netD.apply(weights_init)
netG.apply(weights_init)
else:
raise ValueError('method argument can be either sourceonly or LSD')
if resume:
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']
else:
vgg16 = torchfcn.models.VGG16(pretrained=True)
model.copy_params_from_vgg16(vgg16)
if cuda:
model = model.cuda()
if args.method == 'LSD':
netD = netD.cuda()
netG = netG.cuda()
# Defining optimizer
if args.optimizer == 'SGD':
optim = torch.optim.SGD([
{
'params': get_parameters(model, bias=False)
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2,
'weight_decay': args.weight_decay
},
],
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer == 'Adam':
optim = torch.optim.Adam([
{
'params': get_parameters(model, bias=False)
},
{
'params': get_parameters(model, bias=True),
'lr': args.lr * 2
},
],
lr=args.lr,
betas=(args.beta1, 0.999))
else:
raise ValueError('Invalid optmizer argument. Has to be SGD or Adam')
if args.method == 'LSD':
optimD = torch.optim.Adam(netD.parameters(),
lr=0.0001,
betas=(0.7, 0.999))
optimG = torch.optim.Adam(netG.parameters(),
lr=0.0001,
betas=(0.7, 0.999))
if resume:
optim.load_state_dict(checkpoint['optim_state_dict'])
# Defining trainer object, and start training
if args.method == 'sourceonly':
trainer = torchfcn.Trainer_sourceonly(
cuda=cuda,
model=model,
optimizer=optim,
train_loader=train_loader,
target_loader=target_loader,
val_loader=val_loader,
out=out,
max_iter=args.num_iters,
interval_validate=args.interval_validate)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
elif args.method == 'LSD':
trainer = Trainer_LSD(cuda=cuda,
model=model,
netD=netD,
netG=netG,
optimizer=optim,
optimizerD=optimD,
optimizerG=optimG,
train_loader=train_loader,
target_loader=target_loader,
l1_weight=args.l1_weight,
adv_weight=args.adv_weight,
c_weight=args.c_weight,
val_loader=val_loader,
out=out,
max_iter=args.num_iters,
interval_validate=args.interval_validate,
image_size=image_size)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
