def train_fcn(train_img_path, pths_path, batch_size, lr, num_workers,
epoch_iter, interval):
#file_num = 1056 #暂定
if (not os.path.exists(pths_path)):
os.makedirs(pths_path)
trainset = fcn_dataset(train_img_path)
train_loader = data.DataLoader(trainset, batch_size=batch_size, \
shuffle=True, num_workers=num_workers, drop_last=False)
#criterion = cross_entropy2d()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#device = torch.device("cpu")
vgg_model = VGGNet(pretrained=False, requires_grad=True, remove_fc=True)
fcn_model = FCN8s(pretrained_net=vgg_model, n_class=2)
vgg_model.to(device)
fcn_model.to(device)
data_parallel = False
#model.to(device)
optimizer = torch.optim.Adam(fcn_model.parameters(), lr=lr)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=[epoch_iter // 2],
gamma=0.1)
#criterion = BinaryDiceLoss()
for epoch in range(epoch_iter):
fcn_model.train()
epoch_loss = 0
epoch_time = time.time()
for i, (img, mask) in enumerate(train_loader):
start_time = time.time()
img, mask = img.to(device), mask.to(device)
output = fcn_model(img)
#loss = nn.BCEWithLogitsLoss(output, mask)
loss = cross_entropy2d(output, mask)
#loss = get_dice_loss(output, mask)
#loss = criterion(output, mask)
#loss /= len(img)
epoch_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
print('Epoch is [{}/{}], time consumption is {:.8f}, batch loss is {:.8f}'.format(\
epoch+1, epoch_iter, time.time()-start_time, loss.item()))
print('epoch_loss is {:.8f}, epoch_time is {:.8f}'.format(
epoch_loss / len(img),
time.time() - epoch_time))
print(time.asctime(time.localtime(time.time())))
print('=' * 50)
if (epoch + 1) % interval == 0:
state_dict = fcn_model.module.state_dict(
) if data_parallel else fcn_model.state_dict()
torch.save(
state_dict,
os.path.join(pths_path,
'model_epoch_{}.pth'.format(epoch + 1)))
def main():
# args = parse_args()
torch.backends.cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
# # if args.seed:
# random.seed(args.seed)
# np.random.seed(args.seed)
# torch.manual_seed(args.seed)
# # if args.gpu:
# torch.cuda.manual_seed_all(args.seed)
seed = 63
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# if args.gpu:
torch.cuda.manual_seed_all(seed)
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
# train_transforms = transforms.Compose([
# transforms.RandomCrop(args['crop_size']),
# transforms.RandomRotation(90),
# transforms.RandomHorizontalFlip(p=0.5),
# transforms.RandomVerticalFlip(p=0.5),
# ])
short_size = int(min(args['input_size']) / 0.875)
# val_transforms = transforms.Compose([
# transforms.Scale(short_size, interpolation=Image.NEAREST),
# # joint_transforms.Scale(short_size),
# transforms.CenterCrop(args['input_size'])
# ])
train_joint_transform = joint_transforms.Compose([
# joint_transforms.Scale(short_size),
joint_transforms.RandomCrop(args['crop_size']),
joint_transforms.RandomHorizontallyFlip(),
joint_transforms.RandomRotate(90)
])
val_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.CenterCrop(args['input_size'])
])
input_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(*mean_std)])
target_transform = extended_transforms.MaskToTensor()
restore_transform = transforms.Compose(
[extended_transforms.DeNormalize(*mean_std),
transforms.ToPILImage()])
visualize = transforms.ToTensor()
train_set = cityscapes.CityScapes('train',
joint_transform=train_joint_transform,
transform=input_transform,
target_transform=target_transform)
# train_set = cityscapes.CityScapes('train', transform=train_transforms)
train_loader = DataLoader(train_set,
batch_size=args['train_batch_size'],
num_workers=8,
shuffle=True)
val_set = cityscapes.CityScapes('val',
joint_transform=val_joint_transform,
transform=input_transform,
target_transform=target_transform)
# val_set = cityscapes.CityScapes('val', transform=val_transforms)
val_loader = DataLoader(val_set,
batch_size=args['val_batch_size'],
num_workers=8,
shuffle=True)
print(len(train_loader), len(val_loader))
# sdf
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
net = FCN8s(pretrained_net=vgg_model,
n_class=cityscapes.num_classes,
dropout_rate=0.4)
# net.apply(init_weights)
criterion = nn.CrossEntropyLoss(ignore_index=cityscapes.ignore_label)
optimizer = optim.Adam(net.parameters(), lr=1e-4)
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
open(
os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(args) + '\n\n')
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, 'min', patience=args['lr_patience'], min_lr=1e-10)
vgg_model = vgg_model.to(device)
net = net.to(device)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
if len(args['snapshot']) == 0:
curr_epoch = 1
args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0
}
else:
print('training resumes from ' + args['snapshot'])
net.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9][:-4])
}
criterion.to(device)
for epoch in range(curr_epoch, args['epoch_num'] + 1):
train(train_loader, net, device, criterion, optimizer, epoch, args)
val_loss = validate(val_loader, net, device, criterion, optimizer,
epoch, args, restore_transform, visualize)
scheduler.step(val_loss)
def main():
torch.backends.cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
net = FCN8s(pretrained_net=vgg_model,
n_class=cityscapes.num_classes,
dropout_rate=0.4)
print('load model ' + args['snapshot'])
vgg_model = vgg_model.to(device)
net = net.to(device)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net)
net.load_state_dict(
torch.load(os.path.join(ckpt_path, args['exp_name'],
args['snapshot'])))
net.eval()
mean_std = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
short_size = int(min(args['input_size']) / 0.875)
val_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(short_size),
joint_transforms.CenterCrop(args['input_size'])
])
test_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(*mean_std)])
target_transform = extended_transforms.MaskToTensor()
restore_transform = transforms.Compose(
[extended_transforms.DeNormalize(*mean_std),
transforms.ToPILImage()])
# test_set = cityscapes.CityScapes('test', transform=test_transform)
test_set = cityscapes.CityScapes('test',
joint_transform=val_joint_transform,
transform=test_transform,
target_transform=target_transform)
test_loader = DataLoader(test_set,
batch_size=1,
num_workers=8,
shuffle=False)
transform = transforms.ToPILImage()
check_mkdir(os.path.join(ckpt_path, args['exp_name'], 'test'))
gts_all, predictions_all = [], []
count = 0
for vi, data in enumerate(test_loader):
# img_name, img = data
img_name, img, gts = data
img_name = img_name[0]
# print(img_name)
img_name = img_name.split('/')[-1]
# img.save(os.path.join(ckpt_path, args['exp_name'], 'test', img_name))
img_transform = restore_transform(img[0])
# img_transform = img_transform.convert('RGB')
img_transform.save(
os.path.join(ckpt_path, args['exp_name'], 'test', img_name))
img_name = img_name.split('_leftImg8bit.png')[0]
# img = Variable(img, volatile=True).cuda()
img, gts = img.to(device), gts.to(device)
output = net(img)
prediction = output.data.max(1)[1].squeeze_(1).squeeze_(
0).cpu().numpy()
prediction_img = cityscapes.colorize_mask(prediction)
# print(type(prediction_img))
prediction_img.save(
os.path.join(ckpt_path, args['exp_name'], 'test',
img_name + '.png'))
# print(ckpt_path, args['exp_name'], 'test', img_name + '.png')
print('%d / %d' % (vi + 1, len(test_loader)))
gts_all.append(gts.data.cpu().numpy())
predictions_all.append(prediction)
# break
# if count == 1:
# break
# count += 1
gts_all = np.concatenate(gts_all)
predictions_all = np.concatenate(prediction)
acc, acc_cls, mean_iou, _ = evaluate(predictions_all, gts_all,
cityscapes.num_classes)
print(
'-----------------------------------------------------------------------------------------------------------'
)
print('[acc %.5f], [acc_cls %.5f], [mean_iu %.5f]' %
(acc, acc_cls, mean_iu))
def main():
args = parse_args()
torch.backends.cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
# Random seed for reproducibility
if args.seed:
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.gpu:
torch.cuda.manual_seed_all(args.seed)
# seed = 63
# random.seed(seed)
# np.random.seed(seed)
# torch.manual_seed(seed)
# # if args.gpu:
# torch.cuda.manual_seed_all(seed)
denoramlize_argument = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
train_transforms = transforms.Compose([
transforms.RandomCrop((args.crop_size, args.crop_size)),
transforms.RandomRotation(90),
transforms.RandomHorizontalFlip(p=0.5),
])
# train_joint_transform = joint_transforms.Compose([
# # joint_transforms.Scale(img_resize_shape),
# joint_transforms.RandomCrop(args['crop_size']),
# joint_transforms.RandomHorizontallyFlip(),
# joint_transforms.RandomRotate(90)
# ])
img_resize_shape = int(min(args.input_size) / 0.8)
# val_transforms = transforms.Compose([
# transforms.Scale(img_resize_shape, interpolation=Image.NEAREST),
# transforms.CenterCrop(args['input_size'])
# ])
val_joint_transform = joint_transforms.Compose([
joint_transforms.Scale(img_resize_shape),
joint_transforms.CenterCrop(args.input_size)
])
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = transforms.Compose([
extended_transforms.DeNormalize(*denoramlize_argument),
transforms.ToPILImage()
])
visualize = transforms.ToTensor()
# train_set = games_data.CityScapes('train', joint_transform=train_joint_transform,
# transform=input_transform, target_transform=target_transform)
train_set = games_data.CityScapes('train', transform=train_transforms)
# train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=8, shuffle=True)
train_loader = DataLoader(train_set,
batch_size=args.training_batch_size,
num_workers=8,
shuffle=True)
val_set = games_data.CityScapes('val',
joint_transform=val_joint_transform,
transform=input_transform,
target_transform=target_transform)
val_loader = DataLoader(val_set,
batch_size=args.val_batch_size,
num_workers=8,
shuffle=True)
print(len(train_loader), len(val_loader))
# sdf
# Load pretrained VGG model
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
# FCN architecture load
model = FCN8s(pretrained_net=vgg_model,
n_class=games_data.num_classes,
dropout_rate=0.4)
# Loss function
criterion = nn.CrossEntropyLoss(ignore_index=games_data.ignore_label)
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=1e-4)
# Create directory for checkpoints
exist_directory(ckpt_path)
exist_directory(os.path.join(ckpt_path, exp_name))
open(
os.path.join(ckpt_path, exp_name,
str(datetime.datetime.now()) + '.txt'),
'w').write(str(args) + '\n\n')
# Learning rate scheduler
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
'min',
patience=args.lr_patience,
min_lr=1e-10)
# Send model to CUDA device
vgg_model = vgg_model.to(device)
model = model.to(device)
# Use if more than 1 GPU
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
if len(args.snapshot) == 0:
curr_epoch = 1
best_args['best_record'] = {
'epoch': 0,
'val_loss': 1e10,
'acc': 0,
'acc_cls': 0,
'mean_iu': 0
}
else:
print('training resumes from ' + args['snapshot'])
model.load_state_dict(
torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'])))
split_snapshot = args['snapshot'].split('_')
curr_epoch = int(split_snapshot[1]) + 1
best_args['best_record'] = {
'epoch': int(split_snapshot[1]),
'val_loss': float(split_snapshot[3]),
'acc': float(split_snapshot[5]),
'acc_cls': float(split_snapshot[7]),
'mean_iu': float(split_snapshot[9][:-4])
}
criterion.to(device)
for epoch in range(curr_epoch, args.epochs + 1):
train(train_loader, model, device, criterion, optimizer, epoch, args)
val_loss = validate(val_loader, model, device, criterion, optimizer,
epoch, args, restore_transform, visualize)
scheduler.step(val_loss)