def main():
global args, logger
args = get_parser().parse_args()
logger = get_logger()
# os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h - 1) % 8 == 0 and (args.crop_w - 1) % 8 == 0
assert args.split in ['train', 'val', 'test']
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
gray_folder = os.path.join(args.save_folder, 'gray')
color_folder = os.path.join(args.save_folder, 'color')
val_transform = transforms.Compose([transforms.ToTensor()])
val_data1 = datasets.SegData(split=args.split, data_root=args.data_root, data_list=args.val_list1, transform=val_transform)
val_loader1 = torch.utils.data.DataLoader(val_data1, batch_size=1, shuffle=False, num_workers=args.workers, pin_memory=True)
colors = np.loadtxt(args.colors_path).astype('uint8')
names = [line.rstrip('\n') for line in open(args.names_path)]
if not args.has_prediction:
if args.net_type == 0:
from pspnet import PSPNet
model = PSPNet(backbone = args.backbone, layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, use_softmax=True, use_aux=False, pretrained=False, syncbn=False).cuda()
elif args.net_type in [1, 2, 3]:
from pspnet_div4 import PSPNet
model = PSPNet(backbone = args.backbone, layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, use_softmax=True, use_aux=False, pretrained=False, syncbn=False, net_type=args.net_type).cuda()
logger.info(model)
model = torch.nn.DataParallel(model).cuda()
cudnn.enabled = True
cudnn.benchmark = True
if os.path.isfile(args.model_path):
logger.info("=> loading checkpoint '{}'".format(args.model_path))
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['state_dict'], strict=False)
logger.info("=> loaded checkpoint '{}'".format(args.model_path))
else:
raise RuntimeError("=> no checkpoint found at '{}'".format(args.model_path))
cv2.setNumThreads(0)
mIoUs = []
mAccs = []
allAccs = []
validate(val_loader1, val_data1.data_list, model, args.classes, mean, std, args.base_size1, args.crop_h, args.crop_w, args.scales, gray_folder, color_folder, colors)
if args.split != 'test':
mIoU, mAcc, allAcc = cal_acc(val_data1.data_list, gray_folder, args.classes, names)
mIoUs.append(mIoU)
mAccs.append(mAcc)
allAccs.append(allAcc)
def main():
global args, logger
args = get_parser().parse_args()
logger = get_logger()
# os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h - 1) % 8 == 0 and (args.crop_w - 1) % 8 == 0
assert args.split in ['train', 'val', 'test']
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
train_transform = transforms.Compose([
transforms.RandScale([0.5, 2]),
transforms.RandRotate([-10, 10], padding=mean, ignore_label=args.ignore_label),
transforms.RandomGaussianBlur(),
transforms.RandomHorizontalFlip(),
transforms.Crop([args.crop_h, args.crop_w], crop_type='rand', padding=mean, ignore_label=args.ignore_label),
transforms.ToTensor()])
val_transform = transforms.Compose([transforms.Crop([args.crop_h, args.crop_w], crop_type='center', padding=mean, ignore_label=args.ignore_label),
transforms.ToTensor()])
val_data1 = datasets.SegData(split='train', data_root=args.data_root, data_list=args.val_list1, transform=val_transform)
val_loader1 = torch.utils.data.DataLoader(val_data1, batch_size=1, shuffle=False, num_workers=args.workers, pin_memory=True)
from pspnet import PSPNet
model = PSPNet(backbone = args.backbone, layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, use_softmax=False, use_aux=False, pretrained=False, syncbn=False).cuda()
logger.info(model)
# model = torch.nn.DataParallel(model).cuda()
model = model.cuda()
cudnn.enabled = True
cudnn.benchmark = True
if os.path.isfile(args.model_path):
logger.info("=> loading checkpoint '{}'".format(args.model_path))
checkpoint = torch.load(args.model_path)
# model.load_state_dict(checkpoint['state_dict'], strict=False)
# logger.info("=> loaded checkpoint '{}'".format(args.model_path))
pretrained_dict = {k.replace('module.',''): v for k, v in checkpoint['state_dict'].items()}
dict1 = model.state_dict()
model.load_state_dict(pretrained_dict, strict=False)
else:
raise RuntimeError("=> no checkpoint found at '{}'".format(args.model_path))
cv2.setNumThreads(0)
validate(val_loader1, val_data1.data_list, model, args.classes, mean, std, args.base_size1, args.crop_h, args.crop_w, args.scales)
def build_network(snapshot):
epoch = 0
net = PSPNet()
net = nn.DataParallel(net)
if snapshot is not None:
_, epoch = os.path.basename(snapshot).split('_')
epoch = int(epoch)
net.load_state_dict(torch.load(snapshot))
logging.info("Snapshot for epoch {} loaded from {}".format(
epoch, snapshot))
net = net.cuda()
return net, epoch
def main():
global args, logger
args = get_parser().parse_args()
logger = get_logger()
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h - 1) % 8 == 0 and (args.crop_w - 1) % 8 == 0
assert args.split in ['train', 'val', 'test']
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
gray_folder = os.path.join(args.save_folder, 'gray')
color_folder = os.path.join(args.save_folder, 'color')
val_transform = transforms.Compose([transforms.ToTensor()])
val_data = datasets.SegData(split=args.split, data_root=args.data_root, data_list=args.val_list, transform=val_transform)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=1, shuffle=False, num_workers=args.workers, pin_memory=True)
colors = scio.loadmat(args.colors_path)['colors']
colors = np.uint8(colors)
names = scio.loadmat(args.names_path)['names']
names = [names[i, 0][0] for i in range(0, args.classes)]
if not args.has_prediction:
model = PSPNet(layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, use_softmax=False, use_aux=False, pretrained=False).cuda()
logger.info(model)
model = torch.nn.DataParallel(model).cuda()
cudnn.enabled = True
cudnn.benchmark = True
if os.path.isfile(args.model_path):
logger.info("=> loading checkpoint '{}'".format(args.model_path))
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['state_dict'], strict=False)
logger.info("=> loaded checkpoint '{}'".format(args.model_path))
else:
raise RuntimeError("=> no checkpoint found at '{}'".format(args.model_path))
validate(val_loader, val_data.data_list, model, args.classes, mean, std, args.base_size, args.crop_h, args.crop_w, args.scales, gray_folder, color_folder, colors)
if args.split is not 'test':
cal_acc(val_data.data_list, gray_folder, args.classes, names)
def build_network(snapshot, backend):
epoch = 0
backend = backend.lower()
net = None
if backend.startswith('resnet'):
net = PSPNet(sizes=(1, 2, 3, 6),
psp_size=2048,
deep_features_size=1024,
backend=backend)
net = nn.DataParallel(net)
if snapshot is not None:
_, epoch = os.path.basename(snapshot).split('_')
epoch = int(epoch)
net.load_state_dict(torch.load(snapshot))
logging.info("Snapshot for epoch {} loaded from {}".format(
epoch, snapshot))
net = net.cuda()
return net, epoch
def main():
global args, logger, writer
args = get_parser().parse_args()
logger = get_logger()
writer = SummaryWriter(args.save_path)
# os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
if args.dist:
dist_init(args.port, backend=args.backend)
if len(args.gpu) == 1:
args.syncbn = False
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h - 1) % 8 == 0 and (args.crop_w - 1) % 8 == 0
world_size = 1
rank = 0
if args.dist:
rank = dist.get_rank()
world_size = dist.get_world_size()
if rank == 0:
logger.info('dist:{}'.format(args.dist))
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
# rank = dist.get_rank()
if args.bn_group > 1:
args.syncbn = True
bn_sync_stats = True
bn_group_comm = simple_group_split(world_size, rank,
world_size // args.bn_group)
else:
args.syncbn = False
bn_sync_stats = False
bn_group_comm = None
model = PSPNet(layers=args.layers,
classes=args.classes,
zoom_factor=args.zoom_factor,
syncbn=args.syncbn,
group_size=args.bn_group,
group=bn_group_comm,
sync_stats=bn_sync_stats)
if rank == 0:
logger.info(model)
# optimizer = torch.optim.SGD(model.parameters(), args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
# newly introduced layer with lr x10
optimizer = torch.optim.SGD([{
'params': model.layer0.parameters()
}, {
'params': model.layer1.parameters()
}, {
'params': model.layer2.parameters()
}, {
'params': model.layer3.parameters()
}, {
'params': model.layer4.parameters()
}, {
'params': model.ppm.parameters(),
'lr': args.base_lr * 10
}, {
'params': model.cls.parameters(),
'lr': args.base_lr * 10
}, {
'params': model.aux.parameters(),
'lr': args.base_lr * 10
}],
lr=args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# model = torch.nn.DataParallel(model).cuda()
# if args.syncbn:
# from lib.syncbn import patch_replication_callback
# patch_replication_callback(model)
model = model.cuda()
cudnn.enabled = True
cudnn.benchmark = True
criterion = nn.NLLLoss(ignore_index=args.ignore_label)
if args.weight:
def map_func(storage, location):
return storage.cuda()
if os.path.isfile(args.weight):
logger.info("=> loading weight '{}'".format(args.weight))
checkpoint = torch.load(args.weight, map_location=map_func)
model.load_state_dict(checkpoint['state_dict'])
logger.info("=> loaded weight '{}'".format(args.weight))
else:
logger.info("=> no weight found at '{}'".format(args.weight))
if args.resume:
if os.path.isfile(args.resume):
logger.info("=> loading checkpoint '{}'".format(args.resume))
# checkpoint = torch.load(args.resume)
# args.start_epoch = checkpoint['epoch']
# model.load_state_dict(checkpoint['state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer'])
model, optimizer, args.start_epoch = restore_from(
model, optimizer, args.resume)
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
else:
logger.info("=> no checkpoint found at '{}'".format(args.resume))
if args.dist:
broadcast_params(model)
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
train_transform = transforms.Compose([
transforms.RandScale([args.scale_min, args.scale_max]),
transforms.RandRotate([args.rotate_min, args.rotate_max],
padding=mean,
ignore_label=args.ignore_label),
transforms.RandomGaussianBlur(),
transforms.RandomHorizontalFlip(),
transforms.Crop([args.crop_h, args.crop_w],
crop_type='rand',
padding=mean,
ignore_label=args.ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
train_data = datasets.SegData(split='train',
data_root=args.data_root,
data_list=args.train_list,
transform=train_transform)
train_sampler = None
if args.dist:
train_sampler = DistributedSampler(train_data)
train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.batch_size,
shuffle=False if train_sampler else True,
num_workers=args.workers,
pin_memory=False,
sampler=train_sampler)
if args.evaluate:
val_transform = transforms.Compose([
transforms.Crop([args.crop_h, args.crop_w],
crop_type='center',
padding=mean,
ignore_label=args.ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
val_data = datasets.SegData(split='val',
data_root=args.data_root,
data_list=args.val_list,
transform=val_transform)
val_sampler = None
if args.dist:
val_sampler = DistributedSampler(val_data)
val_loader = torch.utils.data.DataLoader(
val_data,
batch_size=args.batch_size_val,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
sampler=val_sampler)
for epoch in range(args.start_epoch, args.epochs + 1):
loss_train, mIoU_train, mAcc_train, allAcc_train = train(
train_loader, model, criterion, optimizer, epoch, args.zoom_factor,
args.batch_size, args.aux_weight)
writer.add_scalar('loss_train', loss_train.cpu().numpy(), epoch)
writer.add_scalar('mIoU_train', mIoU_train, epoch)
writer.add_scalar('mAcc_train', mAcc_train, epoch)
writer.add_scalar('allAcc_train', allAcc_train, epoch)
# write parameters histogram costs lots of time
# for name, param in model.named_parameters():
# writer.add_histogram(name, param, epoch)
if args.evaluate and rank == 0:
loss_val, mIoU_val, mAcc_val, allAcc_val = validate(
val_loader, model, criterion, args.classes, args.zoom_factor)
writer.add_scalar('loss_val', loss_val.cpu().numpy(), epoch)
writer.add_scalar('mIoU_val', mIoU_val, epoch)
writer.add_scalar('mAcc_val', mAcc_val, epoch)
writer.add_scalar('allAcc_val', allAcc_val, epoch)
if epoch % args.save_step == 0 and (rank == 0):
filename = args.save_path + '/train_epoch_' + str(epoch) + '.pth'
logger.info('Saving checkpoint to: ' + filename)
torch.save(
{
'epoch': epoch,
'state_dict': model.cpu().state_dict(),
'optimizer': optimizer.state_dict()
}, filename)
def main():
global args, logger, writer
args = get_parser().parse_args()
import multiprocessing as mp
if mp.get_start_method(allow_none=True) != 'spawn':
mp.set_start_method('spawn', force=True)
rank, world_size = dist_init(args.port)
logger = get_logger()
writer = SummaryWriter(args.save_path)
#os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
#if len(args.gpu) == 1:
# args.syncbn = False
if rank == 0:
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h - 1) % 8 == 0 and (args.crop_w - 1) % 8 == 0
assert args.net_type in [0, 1, 2, 3]
if args.bn_group == 1:
args.bn_group_comm = None
else:
assert world_size % args.bn_group == 0
args.bn_group_comm = simple_group_split(world_size, rank,
world_size // args.bn_group)
if rank == 0:
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
from pspnet import PSPNet
model = PSPNet(backbone=args.backbone,
layers=args.layers,
classes=args.classes,
zoom_factor=args.zoom_factor,
syncbn=args.syncbn,
group_size=args.bn_group,
group=args.bn_group_comm,
use_softmax=False,
use_aux=False).cuda()
logger.info(model)
# optimizer = torch.optim.SGD(model.parameters(), args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
# newly introduced layer with lr x10
optimizer = torch.optim.SGD(
[{
'params': model.layer0.parameters()
}, {
'params': model.layer1.parameters()
}, {
'params': model.layer2.parameters()
}, {
'params': model.layer3.parameters()
}, {
'params': model.layer4.parameters()
}, {
'params': model.ppm.parameters(),
'lr': args.base_lr * 10
}, {
'params': model.cls.parameters(),
'lr': args.base_lr * 10
}, {
'params': model.result.parameters(),
'lr': args.base_lr * 10
}],
# {'params': model.aux.parameters(), 'lr': args.base_lr * 10}],
lr=args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
#model = torch.nn.DataParallel(model).cuda()
model = DistModule(model)
#if args.syncbn:
# from lib.syncbn import patch_replication_callback
# patch_replication_callback(model)
cudnn.enabled = True
cudnn.benchmark = True
criterion = nn.NLLLoss(ignore_index=args.ignore_label).cuda()
if args.weight:
def map_func(storage, location):
return storage.cuda()
if os.path.isfile(args.weight):
logger.info("=> loading weight '{}'".format(args.weight))
checkpoint = torch.load(args.weight, map_location=map_func)
model.load_state_dict(checkpoint['state_dict'])
logger.info("=> loaded weight '{}'".format(args.weight))
else:
logger.info("=> no weight found at '{}'".format(args.weight))
if args.resume:
load_state(args.resume, model, optimizer)
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
normalize = Normalize()
train_data = voc12.data.VOC12ClsDataset(
args.train_list,
voc12_root=args.voc12_root,
transform=transforms.Compose([
imutils.RandomResizeLong(400, 512),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.3,
contrast=0.3,
saturation=0.3,
hue=0.1), np.asarray, normalize,
imutils.RandomCrop(args.crop_size), imutils.HWC_to_CHW,
torch.from_numpy
]))
train_sampler = DistributedSampler(train_data)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
sampler=train_sampler)
for epoch in range(args.start_epoch, args.epochs + 1):
loss_train = train(train_loader, model, criterion, optimizer, epoch,
args.zoom_factor, args.batch_size, args.aux_weight)
if rank == 0:
writer.add_scalar('loss_train', loss_train, epoch)
# write parameters histogram costs lots of time
# for name, param in model.named_parameters():
# writer.add_histogram(name, param, epoch)
if epoch % args.save_step == 0 and rank == 0:
filename = args.save_path + '/train_epoch_' + str(epoch) + '.pth'
logger.info('Saving checkpoint to: ' + filename)
torch.save(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, filename)
def main():
global args, logger
args = get_parser().parse_args()
logger = get_logger()
# os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
logger.info(args)
logger.info("=> creating model ...")
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
val_transform = transforms.Compose([
transforms.Resize(size=(256, 256)),
transforms.Crop([args.crop_h, args.crop_w],
crop_type='center',
padding=mean,
ignore_label=args.ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
val_data1 = datasets.SegData(split=args.split,
data_root=args.data_root,
data_list=args.val_list1,
transform=val_transform)
val_loader1 = torch.utils.data.DataLoader(val_data1,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
model_pfr = PFR().cuda()
model_pfr = torch.nn.DataParallel(model_pfr)
model_prp = PRP().cuda()
model_prp = torch.nn.DataParallel(model_prp)
from pspnet import PSPNet
model = PSPNet(backbone=args.backbone,
layers=args.layers,
classes=args.classes,
zoom_factor=args.zoom_factor,
use_softmax=True,
pretrained=False,
syncbn=False).cuda()
logger.info(model)
model = torch.nn.DataParallel(model).cuda()
cudnn.enabled = True
cudnn.benchmark = True
if os.path.isfile(args.model_path):
logger.info("=> loading checkpoint '{}'".format(args.model_path))
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['state_dict'], strict=False)
logger.info("=> loaded checkpoint '{}'".format(args.model_path))
else:
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.model_path))
checkpoint_pfr = torch.load(args.model_path.replace('.pth', '_pfr.pth'))
checkpoint_prp = torch.load(args.model_path.replace('.pth', '_prp.pth'))
model_pfr.load_state_dict(checkpoint_pfr['state_dict'], strict=False)
model_prp.load_state_dict(checkpoint_prp['state_dict'], strict=False)
cv2.setNumThreads(0)
validate(val_loader1, val_data1.data_list, model, model_pfr, model_prp)
def main():
global args, logger
args = get_parser().parse_args()
logger = get_logger()
# os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h - 1) % 8 == 0 and (args.crop_w - 1) % 8 == 0
assert args.split in ['train', 'val', 'test']
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
normalize = Normalize()
infer_dataset = voc12.data.VOC12ClsDataset(
args.train_list,
voc12_root=args.voc12_root,
transform=transforms.Compose([
np.asarray,
imutils.RandomCrop(441),
# normalize,
imutils.HWC_to_CHW
]))
val_loader1 = DataLoader(infer_dataset,
batch_size=1,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
from pspnet import PSPNet
model = PSPNet(backbone=args.backbone,
layers=args.layers,
classes=args.classes,
zoom_factor=args.zoom_factor,
use_softmax=False,
use_aux=False,
pretrained=False,
syncbn=False).cuda()
logger.info(model)
# model = torch.nn.DataParallel(model).cuda()
model = model.cuda()
cudnn.enabled = True
cudnn.benchmark = True
if os.path.isfile(args.model_path):
logger.info("=> loading checkpoint '{}'".format(args.model_path))
checkpoint = torch.load(args.model_path)
# model.load_state_dict(checkpoint['state_dict'], strict=False)
# logger.info("=> loaded checkpoint '{}'".format(args.model_path))
pretrained_dict = {
k.replace('module.', ''): v
for k, v in checkpoint['state_dict'].items()
}
dict1 = model.state_dict()
model.load_state_dict(pretrained_dict, strict=False)
else:
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.model_path))
cv2.setNumThreads(0)
validate(val_loader1, model, args.classes, mean, std, args.base_size1)
def main():
global args, logger, writer
args = get_parser().parse_args()
import multiprocessing as mp
if mp.get_start_method(allow_none=True) != 'spawn':
mp.set_start_method('spawn', force=True)
rank, world_size = dist_init(args.port)
logger = get_logger()
writer = SummaryWriter(args.save_path)
#os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
#if len(args.gpu) == 1:
# args.syncbn = False
if rank == 0:
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h-1) % 8 == 0 and (args.crop_w-1) % 8 == 0
assert args.net_type in [0, 1, 2, 3]
if args.bn_group == 1:
args.bn_group_comm = None
else:
assert world_size % args.bn_group == 0
args.bn_group_comm = simple_group_split(world_size, rank, world_size // args.bn_group)
if rank == 0:
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
if args.net_type == 0:
from pspnet import PSPNet
model = PSPNet(backbone=args.backbone, layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, syncbn=args.syncbn, group_size=args.bn_group, group=args.bn_group_comm).cuda()
elif args.net_type in [1, 2, 3]:
from pspnet_div4 import PSPNet
model = PSPNet(backbone=args.backbone, layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, syncbn=args.syncbn, group_size=args.bn_group, group=args.bn_group_comm, net_type=args.net_type).cuda()
logger.info(model)
# optimizer = torch.optim.SGD(model.parameters(), args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
# newly introduced layer with lr x10
if args.net_type == 0:
optimizer = torch.optim.SGD(
[{'params': model.layer0.parameters()},
{'params': model.layer1.parameters()},
{'params': model.layer2.parameters()},
{'params': model.layer3.parameters()},
{'params': model.layer4.parameters()},
{'params': model.ppm.parameters(), 'lr': args.base_lr * 10},
{'params': model.conv6.parameters(), 'lr': args.base_lr * 10},
{'params': model.conv1_1x1.parameters(), 'lr': args.base_lr * 10},
{'params': model.cls.parameters(), 'lr': args.base_lr * 10},
{'params': model.aux.parameters(), 'lr': args.base_lr * 10}],
lr=args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
elif args.net_type == 1:
optimizer = torch.optim.SGD(
[{'params': model.layer0.parameters()},
{'params': model.layer1.parameters()},
{'params': model.layer2.parameters()},
{'params': model.layer3.parameters()},
{'params': model.layer4.parameters()},
{'params': model.layer4_p.parameters()},
{'params': model.ppm.parameters(), 'lr': args.base_lr * 10},
{'params': model.ppm_p.parameters(), 'lr': args.base_lr * 10},
{'params': model.cls.parameters(), 'lr': args.base_lr * 10},
{'params': model.cls_p.parameters(), 'lr': args.base_lr * 10},
{'params': model.aux.parameters(), 'lr': args.base_lr * 10}],
lr=args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
elif args.net_type == 2:
optimizer = torch.optim.SGD(
[{'params': model.layer0.parameters()},
{'params': model.layer1.parameters()},
{'params': model.layer2.parameters()},
{'params': model.layer3.parameters()},
{'params': model.layer4.parameters()},
{'params': model.layer4_p.parameters()},
{'params': model.ppm.parameters(), 'lr': args.base_lr * 10},
{'params': model.ppm_p.parameters(), 'lr': args.base_lr * 10},
{'params': model.cls.parameters(), 'lr': args.base_lr * 10},
{'params': model.cls_p.parameters(), 'lr': args.base_lr * 10},
{'params': model.att.parameters(), 'lr': args.base_lr * 10},
{'params': model.att_p.parameters(), 'lr': args.base_lr * 10},
{'params': model.aux.parameters(), 'lr': args.base_lr * 10}],
lr=args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
elif args.net_type == 3:
optimizer = torch.optim.SGD(
[{'params': model.layer0.parameters()},
{'params': model.layer1.parameters()},
{'params': model.layer2.parameters()},
{'params': model.layer3.parameters()},
{'params': model.layer4.parameters()},
{'params': model.layer4_p.parameters()},
{'params': model.ppm.parameters(), 'lr': args.base_lr * 10},
{'params': model.ppm_p.parameters(), 'lr': args.base_lr * 10},
{'params': model.cls.parameters(), 'lr': args.base_lr * 10},
{'params': model.cls_p.parameters(), 'lr': args.base_lr * 10},
{'params': model.att.parameters(), 'lr': args.base_lr * 10},
{'params': model.aux.parameters(), 'lr': args.base_lr * 10}],
lr=args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
fcw = V11RFCN()
fcw_model = torch.load('checkpoint_e8.pth')['state_dict']
fcw_dict = fcw.state_dict()
pretrained_fcw = {k: v for k, v in fcw_model.items() if k in fcw_dict}
fcw_dict.update(pretrained_fcw)
fcw.load_state_dict(fcw_dict)
#fcw = DistModule(fcw)
#print(fcw)
fcw = fcw.cuda()
#model = torch.nn.DataParallel(model).cuda()
model = DistModule(model)
#if args.syncbn:
# from lib.syncbn import patch_replication_callback
# patch_replication_callback(model)
cudnn.enabled = True
cudnn.benchmark = True
criterion = nn.NLLLoss(ignore_index=args.ignore_label).cuda()
if args.weight:
def map_func(storage, location):
return storage.cuda()
if os.path.isfile(args.weight):
logger.info("=> loading weight '{}'".format(args.weight))
checkpoint = torch.load(args.weight, map_location=map_func)['state_dict']
checkpoint = {k: v for k, v in checkpoint.items() if 'ppm' not in k}
model_dict = model.state_dict()
model_dict.update(checkpoint)
model.load_state_dict(model_dict)
logger.info("=> loaded weight '{}'".format(args.weight))
else:
logger.info("=> no weight found at '{}'".format(args.weight))
if args.resume:
load_state(args.resume, model, optimizer)
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
train_transform = transforms.Compose([
transforms.RandScale([args.scale_min, args.scale_max]),
#transforms.RandRotate([args.rotate_min, args.rotate_max], padding=mean, ignore_label=args.ignore_label),
transforms.RandomGaussianBlur(),
transforms.RandomHorizontalFlip(),
transforms.Crop([args.crop_h, args.crop_w], crop_type='rand', padding=mean, ignore_label=args.ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)])
train_data = datasets.SegData(split='train', data_root=args.data_root, data_list=args.train_list, transform=train_transform)
train_sampler = DistributedSampler(train_data)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=False, sampler=train_sampler)
if args.evaluate:
val_transform = transforms.Compose([
transforms.Crop([args.crop_h, args.crop_w], crop_type='center', padding=mean, ignore_label=args.ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)])
val_data = datasets.SegData(split='val', data_root=args.data_root, data_list=args.val_list, transform=val_transform)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=args.batch_size_val, shuffle=False, num_workers=args.workers, pin_memory=True)
for epoch in range(args.start_epoch, args.epochs + 1):
loss_train, mIoU_train, mAcc_train, allAcc_train = train(train_loader, model, criterion, optimizer, epoch, args.zoom_factor, args.batch_size, args.aux_weight, fcw)
if rank == 0:
writer.add_scalar('loss_train', loss_train, epoch)
writer.add_scalar('mIoU_train', mIoU_train, epoch)
writer.add_scalar('mAcc_train', mAcc_train, epoch)
writer.add_scalar('allAcc_train', allAcc_train, epoch)
# write parameters histogram costs lots of time
# for name, param in model.named_parameters():
# writer.add_histogram(name, param, epoch)
if epoch % args.save_step == 0 and rank == 0:
filename = args.save_path + '/train_epoch_' + str(epoch) + '.pth'
logger.info('Saving checkpoint to: ' + filename)
torch.save({'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}, filename)
#if epoch / args.save_step > 2:
# deletename = args.save_path + '/train_epoch_' + str(epoch - args.save_step*2) + '.pth'
# os.remove(deletename)
if args.evaluate:
loss_val, mIoU_val, mAcc_val, allAcc_val = validate(val_loader, model, criterion, args.classes, args.zoom_factor)
writer.add_scalar('loss_val', loss_val, epoch)
writer.add_scalar('mIoU_val', mIoU_val, epoch)
writer.add_scalar('mAcc_val', mAcc_val, epoch)
writer.add_scalar('allAcc_val', allAcc_val, epoch)
parser.add_argument("--out_cam_pred", default=None, type=str)
args = parser.parse_args()
from pspnet import PSPNet
model = PSPNet(backbone = 'resnet', layers=50, classes=20, zoom_factor=1, pretrained=False, syncbn=False).cuda()
checkpoint = torch.load('exp/drivable/res101_psp_coarse/model/train_epoch_14.pth')
pretrained_dict = {k.replace('module.',''): v for k, v in checkpoint['state_dict'].items()}
dict1 = model.state_dict()
print (dict1.keys(), pretrained_dict.keys())
for item in dict1:
if item not in pretrained_dict.keys():
print(item,'nbnmbkjhiuguig~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~`')
model.load_state_dict(pretrained_dict, strict=False)
model.eval()
model.cuda()
print(model)
normalize = Normalize()
infer_dataset = voc12.data.VOC12ClsDatasetMSF(args.infer_list, voc12_root=args.voc12_root,
scales=(1, 0.5, 1.5, 2.0),
inter_transform=torchvision.transforms.Compose(
[np.asarray,
normalize,
imutils.HWC_to_CHW]))
infer_data_loader = DataLoader(infer_dataset, shuffle=False, num_workers=args.num_workers, pin_memory=True)
print('data ready')
n_gpus = torch.cuda.device_count()
def load_model(filepath):
net = PSPNet(pretrained=False)
net.cpu()
net.load_state_dict(torch.load(filepath, map_location='cpu'))
return net
################################# MAIN #################################
if len(sys.argv) != 3:
print("Usage:\n\t python3 {} IMAGE_FILE MODEL_PATH".format(sys.argv[0]))
image_path = sys.argv[1]
model_path = sys.argv[2]
assert os.path.isfile(image_path), "File {} does not exist...".format(
image_path)
assert os.path.isfile(model_path), "Model {} does not exist...".format(
model_path)
model = PSPNet(n_classes=6)
model.load_state_dict(torch.load(model_path))
model.eval()
image_np = cv2.imread(image_path)
image_np = shortside_resize(image_np)
image = image_np / 127.5 - 1
image = np.transpose(image, (2, 0, 1))
cls_map = [
'background',
'person',
'bird',
'car',
'cat',
'plane',
# cv2.imwrite(os.path.join(output_path,image_file+'.jpg'),image)
# colorized_mask.save(os.path.join(output_path, image_file+'.png'))
if __name__ == '__main__':
path = './keyboard.pth' #导入的模型文件必须与当前文件在同一目录下
# path='./seg_hand.pth'
checkpoint = torch.load(path)
if isinstance(checkpoint, dict) and 'state_dict' in checkpoint.keys():
checkpoint = checkpoint['state_dict']
if 'module' in list(checkpoint.keys())[0] and not isinstance(
model, torch.nn.DataParallel):
model = torch.nn.DataParallel(model)
model.load_state_dict(checkpoint)
model.to(device)
model.eval()
paramerters = sum(x.numel() for x in model.parameters())
#---model->51M
print("models have {} M paramerters in total".format(paramerters / 1e6))
img_paths = './mask_imgs'
img_paths = [
os.path.join(img_paths, x) for x in os.listdir(img_paths)
if x.endswith('.png')
]
# for img_path in img_paths:
# if not os.path.basename(img_path)=='0000.png':continue
# img=cv2.imread('./keyboard_images/0015.jpg')
def main():
global args, best_record
args = parser.parse_args()
if args.augment:
transform_train = joint_transforms.Compose([
joint_transforms.FreeScale((512, 512)),
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomVerticallyFlip(),
joint_transforms.Rotate(90),
])
transform_val = joint_transforms.Compose(
[joint_transforms.FreeScale((512, 512))])
else:
transform_train = None
dataset_train = dataset.PRCVData('train', args.data_root,
args.label_train_list, transform_train)
dataloader_train = data.DataLoader(dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=8)
dataset_val = dataset.PRCVData('val', args.data_root, args.label_val_list,
transform_val)
dataloader_val = data.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=None,
num_workers=8)
model = PSPNet(num_classes=args.num_class)
print('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
if args.num_class != 21:
for i in saved_state_dict:
#Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
# print i_parts
if i_parts[0] != 'fc':
new_params[i] = saved_state_dict[i]
model.load_state_dict(new_params)
model = model.cuda()
model = torch.nn.DataParallel(model)
cudnn.benchmark = True
# optionally resume from a checkpoint
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']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# define loss function (criterion) and pptimizer
criterion = torch.nn.CrossEntropyLoss(ignore_index=255).cuda()
optimizer = torch.optim.SGD([{
'params': get_1x_lr_params(model),
'lr': args.learning_rate
}, {
'params': get_10x_lr_params(model),
'lr': 10 * args.learning_rate
}],
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(dataloader_train, model, criterion, optimizer, epoch)
# evaluate on validation set
acc, mean_iou, val_loss = validate(dataloader_val, model, criterion,
args.result_pth, epoch)
is_best = mean_iou > best_record['miou']
if is_best:
best_record['epoch'] = epoch
best_record['val_loss'] = val_loss.avg
best_record['acc'] = acc
best_record['miou'] = mean_iou
save_checkpoint(
{
'epoch': epoch + 1,
'val_loss': val_loss.avg,
'accuracy': acc,
'miou': mean_iou,
'state_dict': model.state_dict(),
}, is_best)
print(
'------------------------------------------------------------------------------------------------------'
)
print('[epoch: %d], [val_loss: %5f], [acc: %.5f], [miou: %.5f]' %
(epoch, val_loss.avg, acc, mean_iou))
print(
'best record: [epoch: {epoch}], [val_loss: {val_loss:.5f}], [acc: {acc:.5f}], [miou: {miou:.5f}]'
.format(**best_record))
print(
'------------------------------------------------------------------------------------------------------'
)
def main():
global args, logger
args = get_parser().parse_args()
logger = get_logger()
# os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
logger.info(args)
assert args.classes > 1
assert args.zoom_factor in [1, 2, 4, 8]
assert (args.crop_h - 1) % 8 == 0 and (args.crop_w - 1) % 8 == 0
assert args.split in ['train', 'val', 'test']
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
val_transform = transforms.Compose([
transforms.Resize((args.crop_h, args.crop_w)),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)
])
val_data1 = datasets.SegData(split=args.split,
data_root=args.data_root,
data_list=args.val_list1,
transform=val_transform)
val_loader1 = torch.utils.data.DataLoader(val_data1,
batch_size=1,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
from pspnet import PSPNet
model = PSPNet(backbone=args.backbone,
layers=args.layers,
classes=args.classes,
zoom_factor=args.zoom_factor,
use_softmax=False,
use_aux=False,
pretrained=False,
syncbn=False).cuda()
logger.info(model)
model = torch.nn.DataParallel(model).cuda()
cudnn.enabled = True
cudnn.benchmark = True
if os.path.isfile(args.model_path):
logger.info("=> loading checkpoint '{}'".format(args.model_path))
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['state_dict'], strict=False)
logger.info("=> loaded checkpoint '{}'".format(args.model_path))
else:
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.model_path))
cv2.setNumThreads(0)
fff = open('/mnt/lustre/share/dingmingyu/cityscapes/instance_list_new.txt'
).readlines()
flag = []
for i, line in enumerate(fff):
if i > 100:
break
img = line.strip().split()[0]
img = cv2.imread(img)
if 'ignore' in line.strip().split()[1]:
gt = cv2.imread(line.strip().split()[2])
flag.append(1)
else:
flag.append(0)
gt = cv2.imread(line.strip().split()[1])
cv2.imwrite('result/result_%d_gt.png' % i, gt)
cv2.imwrite('result/result_%d_ori.png' % i, img)
validate(val_loader1, val_data1.data_list, model, args.classes, mean, std,
args.base_size1, args.crop_h, args.crop_w, flag)
print("Eval Process Starting...")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
val_loader = DataLoader(TinySegData(phase='val'),
batch_size=1,
shuffle=False,
num_workers=0)
print("data loading finished...")
model = PSPNet(n_classes=6).to(device)
criterion = torch.nn.CrossEntropyLoss()
mkdirs = lambda x: os.makedirs(x, exist_ok=True)
model.load_state_dict(torch.load(
"./ckpt_seg/epoch_79_iou0.85.pth")) # the model storing path
# eval
model.eval()
for j, (images, seg_gts, rets) in enumerate(val_loader):
if j % 100 == 0:
print('{} sets finished...'.format(j))
# load data to device
images = images.to(device)
seg_gts = seg_gts.to(device)
# get prediction
seg_logit = model(images)
loss_seg = criterion(seg_logit, seg_gts.long())
seg_preds = torch.argmax(seg_logit, dim=1)
def main():
global args, logger, writer
args = get_parser().parse_args()
import multiprocessing as mp
if mp.get_start_method(allow_none=True) != 'spawn':
mp.set_start_method('spawn', force=True)
rank, world_size = dist_init(args.port)
logger = get_logger()
writer = SummaryWriter(args.save_path)
#os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in args.gpu)
#if len(args.gpu) == 1:
# args.syncbn = False
if rank == 0:
logger.info(args)
if args.bn_group == 1:
args.bn_group_comm = None
else:
assert world_size % args.bn_group == 0
args.bn_group_comm = simple_group_split(world_size, rank, world_size // args.bn_group)
if rank == 0:
logger.info("=> creating model ...")
logger.info("Classes: {}".format(args.classes))
from pspnet import PSPNet
model = PSPNet(backbone=args.backbone, layers=args.layers, classes=args.classes, zoom_factor=args.zoom_factor, syncbn=args.syncbn, group_size=args.bn_group, group=args.bn_group_comm).cuda()
logger.info(model)
model_ppm = PPM().cuda()
# optimizer = torch.optim.SGD(model.parameters(), args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
# newly introduced layer with lr x10
optimizer = torch.optim.SGD(
[{'params': model.layer0.parameters()},
{'params': model.layer1.parameters()},
{'params': model.layer2.parameters()},
{'params': model.layer3.parameters()},
{'params': model.layer4_ICR.parameters()},
{'params': model.layer4_PFR.parameters()},
{'params': model.layer4_PRP.parameters()},
{'params': model_ppm.cls_trans.parameters(), 'lr': args.base_lr * 10},
{'params': model_ppm.cls_quat.parameters(), 'lr': args.base_lr * 10}
],
lr=args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay)
#model = torch.nn.DataParallel(model).cuda()
model = DistModule(model)
model_ppm = DistModule(model_ppm)
cudnn.enabled = True
cudnn.benchmark = True
criterion = nn.L1Loss().cuda()
if args.weight:
def map_func(storage, location):
return storage.cuda()
if os.path.isfile(args.weight):
logger.info("=> loading weight '{}'".format(args.weight))
checkpoint = torch.load(args.weight, map_location=map_func)
model.load_state_dict(checkpoint['state_dict'])
logger.info("=> loaded weight '{}'".format(args.weight))
else:
logger.info("=> no weight found at '{}'".format(args.weight))
if args.resume:
load_state(args.resume, model, model_ppm, optimizer)
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
train_transform = transforms.Compose([
transforms.Resize(size=(256,256)),
#transforms.RandomGaussianBlur(),
transforms.Crop([args.crop_h, args.crop_w], crop_type='rand', padding=mean, ignore_label=args.ignore_label),
transforms.ColorJitter([0.4,0.4,0.4]),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)])
train_data = datasets.SegData(split='train', data_root=args.data_root, data_list=args.train_list, transform=train_transform)
train_sampler = DistributedSampler(train_data)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=False, sampler=train_sampler)
if args.evaluate:
val_transform = transforms.Compose([
transforms.Resize(size=(256,256)),
transforms.Crop([args.crop_h, args.crop_w], crop_type='center', padding=mean, ignore_label=args.ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std)])
val_data = datasets.SegData(split='val', data_root=args.data_root, data_list=args.val_list, transform=val_transform)
val_sampler = DistributedSampler(val_data)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=args.batch_size_val, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler)
for epoch in range(args.start_epoch, args.epochs + 1):
t_loss_train, r_loss_train= train(train_loader, model, model_ppm, criterion, optimizer, epoch, args.zoom_factor, args.batch_size, args.aux_weight)
if rank == 0:
writer.add_scalar('t_loss_train', t_loss_train, epoch)
writer.add_scalar('r_loss_train', r_loss_train, epoch)
# write parameters histogram costs lots of time
# for name, param in model.named_parameters():
# writer.add_histogram(name, param, epoch)
if epoch % args.save_step == 0 and rank == 0:
filename = args.save_path + '/train_epoch_' + str(epoch) + '.pth'
logger.info('Saving checkpoint to: ' + filename)
filename_ppm = args.save_path + '/train_epoch_' + str(epoch) + '_ppm.pth'
torch.save({'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}, filename)
torch.save({'epoch': epoch, 'state_dict': model_ppm.state_dict(), 'optimizer': optimizer.state_dict()}, filename_ppm)
#if epoch / args.save_step > 2:
# deletename = args.save_path + '/train_epoch_' + str(epoch - args.save_step*2) + '.pth'
# os.remove(deletename)
if args.evaluate:
t_loss_val, r_loss_val= validate(val_loader, model, model_ppm, criterion)
writer.add_scalar('t_loss_val', t_loss_val, epoch)
writer.add_scalar('r_loss_val', r_loss_val, epoch)
writer.close()
CLASS_NUM = args.class_num
WEIGHTS = args.weights
COLORS = args.colors
SAMPLES = args.samples
OUTPUTS = args.outputs
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = PSPNet(num_classes=CLASS_NUM,
downsample_factor=16,
pretrained=False,
aux_branch=False).to(device=device)
print('model structure is: ')
print(model)
model.load_state_dict(torch.load(WEIGHTS, map_location=device))
model.eval()
# PyTorch转ONNX
print('=================================')
dummy_input = torch.randn(1, 3, 473, 473).to(device)
torch.onnx.export(model,
dummy_input,
'pspnet.onnx',
dynamic_axes={
'image': {
0: 'B'
},
'outputs': {
0: 'B'
}