def pretrain(dataloader, network, path=None):
class config:
lr = 1e-3
epochs = 100
path = '../checkpoint/pretrained_net.pth'
pretrain_config = config()
if path:
pretrain_config.path = path
network.to(device)
criterion_ = CrossEntropyLoss2d()
optimiser_ = torch.optim.Adam(network.parameters(), pretrain_config.lr)
loss_meter = AverageValueMeter()
for i in range(pretrain_config.epochs):
loss_meter.reset()
for i, (img, mask, weak_mask, _) in tqdm(enumerate(dataloader)):
img, mask = img.to(device), mask.to(device)
optimiser_.zero_grad()
output = network(img)
loss = criterion_(output, mask.squeeze(1))
loss.backward()
optimiser_.step()
loss_meter.add(loss.item())
# import ipdb
# ipdb.set_trace()
print(loss_meter.value()[0])
torch.save(network.state_dict(), pretrain_config.path)
# torch.save(network.parameters(),path)
print('pretrained model saved.')
def __init__(self, neural_network, lowerbound, upperbound):
super(networks, self).__init__()
self.lowbound = lowerbound
self.upbound = upperbound
self.neural_net = neural_network
self.reset()
self.optimiser = torch.optim.Adam(self.neural_net.parameters(), lr=0.001, weight_decay=1e-5)
self.CEloss_criterion = CrossEntropyLoss2d()
self.u_r = 1
self.u_s = 1
def __init__(self, neural_network, lowerbound, upperbound):
super(networks, self).__init__()
self.lowbound = lowerbound
self.upbound = upperbound
self.neural_net = neural_network
self.reset()
self.optimiser = torch.optim.Adam(self.neural_net.parameters(),
lr=0.0001)
self.CEloss_criterion = CrossEntropyLoss2d()
self.u_r = 1000
self.lamda = 0
# self.set_bound=False
self.sigma = 0.05
self.scale = 0.5
def update_theta(f_theta_labeled, f_theta_unlabeled, gamma, s, u, v):
global u_r, u_s, net
global labeled_img, labeled_mask, labeled_weak_mask, unlabeled_img, unlabeled_mask
optimiser = torch.optim.Adam(net.parameters(), lr=1e-3)
criterion = CrossEntropyLoss2d()
for i in xrange(10):
loss_labeled = criterion(f_theta_labeled, labeled_mask.squeeze(1))
loss_unlabeled = u_r * (f_theta_unlabeled - gamma.float() + torch.Tensor(u).float()).norm(p=2) + \
u_s * (f_theta_unlabeled - s.float() + torch.Tensor(v)).norm(p=2)
loss = loss_labeled + loss_unlabeled
optimiser.zero_grad()
loss.backward()
optimiser.step()
f_theta_labeled = net(labeled_img) # shape b,c,w,h
f_theta_unlabeled = net(unlabeled_img)
return f_theta_labeled, f_theta_unlabeled
def train(self):
self.color_transform = Colorization(2)
# Dataset loader for train and test
dataset_train = DataLoader(PCD(
os.path.join(self.args.datadir, 'set{}'.format(self.args.cvset),
'train')),
num_workers=self.args.num_workers,
batch_size=self.args.batch_size,
shuffle=True)
self.dataset_test = PCD(
os.path.join(self.args.datadir, 'set{}'.format(self.args.cvset),
'test'))
self.test_path = os.path.join(self.dn_save, 'test')
if not os.path.exists(self.test_path):
os.makedirs(self.test_path)
# Set loss function, optimizer and learning rate
weight = torch.ones(2)
criterion = CrossEntropyLoss2d(weight.cuda())
optimizer = Adam(self.model.parameters(),
lr=0.0001,
betas=(0.5, 0.999))
lambda1 = lambda icount: (float)(self.args.max_iteration - icount) / (
float)(self.args.max_iteration)
model_lr_scheduler = lr_scheduler.LambdaLR(optimizer,
lr_lambda=lambda1)
fn_loss = os.path.join(self.dn_save, 'loss.csv')
f_loss = open(fn_loss, 'w')
writer = csv.writer(f_loss)
self.writers = SummaryWriter(os.path.join(self.dn_save, 'log'))
# Training loop
icount_loss = []
while self.icount < self.args.max_iteration:
for step, (inputs_train, mask_train) in enumerate(dataset_train):
inputs_train = inputs_train.cuda()
mask_train = mask_train.cuda()
inputs_train = Variable(inputs_train)
mask_train = Variable(mask_train)
outputs_train = self.model(inputs_train)
optimizer.zero_grad()
self.loss = criterion(outputs_train, mask_train[:, 0])
self.loss.backward()
optimizer.step()
self.icount += 1
icount_loss.append(self.loss.item())
writer.writerow([self.icount, self.loss.item()])
if self.args.icount_plot > 0 and self.icount % self.args.icount_plot == 0:
self.test()
average = sum(icount_loss) / len(icount_loss)
print('loss: {0} (icount: {1})'.format(
average, self.icount))
icount_loss.clear()
if self.args.icount_save > 0 and self.icount % self.args.icount_save == 0:
self.checkpoint()
# Call lr_schduler.step() after optimizer.step()
model_lr_scheduler.step()
f_loss.close()
def train(args, model, enc=False):
global best_acc
#TODO: calculate weights by processing dataset histogram (now its being set by hand from the torch values)
#create a loder to run all images and calculate histogram of labels, then create weight array using class balancing
weight = torch.ones(NUM_CLASSES)
if (enc):
weight[0] = 2.3653597831726
weight[1] = 4.4237880706787
weight[2] = 2.9691488742828
weight[3] = 5.3442072868347
weight[4] = 5.2983593940735
weight[5] = 5.2275490760803
weight[6] = 5.4394111633301
weight[7] = 5.3659925460815
weight[8] = 3.4170460700989
weight[9] = 5.2414722442627
weight[10] = 4.7376127243042
weight[11] = 5.2286224365234
weight[12] = 5.455126285553
weight[13] = 4.3019247055054
weight[14] = 5.4264230728149
weight[15] = 5.4331531524658
weight[16] = 5.433765411377
weight[17] = 5.4631009101868
weight[18] = 5.3947434425354
else:
weight[0] = 2.8149201869965
weight[1] = 6.9850029945374
weight[2] = 3.7890393733978
weight[3] = 9.9428062438965
weight[4] = 9.7702074050903
weight[5] = 9.5110931396484
weight[6] = 10.311357498169
weight[7] = 10.026463508606
weight[8] = 4.6323022842407
weight[9] = 9.5608062744141
weight[10] = 7.8698215484619
weight[11] = 9.5168733596802
weight[12] = 10.373730659485
weight[13] = 6.6616044044495
weight[14] = 10.260489463806
weight[15] = 10.287888526917
weight[16] = 10.289801597595
weight[17] = 10.405355453491
weight[18] = 10.138095855713
weight[19] = 0
#loader = DataLoader(VOC12(args.datadir, input_transform, target_transform),
# num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True)
assert os.path.exists(
args.datadir), "Error: datadir (dataset directory) could not be loaded"
co_transform = MyCoTransform(enc, augment=True, height=args.height) #1024)
co_transform_val = MyCoTransform(enc, augment=False,
height=args.height) #1024)
dataset_train = cityscapes(args.datadir, co_transform, 'train')
dataset_val = cityscapes(args.datadir, co_transform_val, 'val')
loader = DataLoader(dataset_train,
num_workers=args.num_workers,
batch_size=args.batch_size,
shuffle=True)
loader_val = DataLoader(dataset_val,
num_workers=args.num_workers,
batch_size=args.batch_size,
shuffle=False)
if args.cuda:
weight = weight.cuda()
criterion = CrossEntropyLoss2d(weight)
print(type(criterion))
savedir = args.dirpath + args.savedir
if (enc):
automated_log_path = savedir + "/automated_log_encoder.txt"
modeltxtpath = savedir + "/model_encoder.txt"
else:
automated_log_path = savedir + "/automated_log.txt"
modeltxtpath = savedir + "/model.txt"
if (not os.path.exists(automated_log_path)
): #dont add first line if it exists
with open(automated_log_path, "a") as myfile:
myfile.write(
"Epoch\t\tTrain-loss\t\tTest-loss\t\tTrain-IoU\t\tTest-IoU\t\tlearningRate"
)
with open(modeltxtpath, "w") as myfile:
myfile.write(str(model))
#TODO: reduce memory in first gpu: https://discuss.pytorch.org/t/multi-gpu-training-memory-usage-in-balance/4163/4 #https://github.com/pytorch/pytorch/issues/1893
"""
#Some optimizer examples:
optimizer = Adam(model.parameters())
if args.model.startswith('FCN'):
optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5)
if args.model.startswith('PSP'):
optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4)
if args.model.startswith('Seg'):
optimizer = SGD(model.parameters(), 1e-3, .9)
if args.model.startswith('E'):
#optimizer = Adam(model.parameters(), 1e-3, .9)
optimizer = Adam(model.parameters(), 5e-4, .9, weight_decay=2e-4)
#5e-4 wd: 2e-4
"""
#optimizer = Adam(model.parameters(), 5e-4, (0.9, 0.999), eps=1e-08, weight_decay=2e-4) ## scheduler 1
optimizer = Adam(model.parameters(),
5e-4, (0.9, 0.999),
eps=1e-08,
weight_decay=1e-4) ## scheduler 2
start_epoch = 1
if args.resume:
#Must load weights, optimizer, epoch and best value.
if enc:
filenameCheckpoint = savedir + '/checkpoint_enc.pth.tar'
else:
filenameCheckpoint = savedir + '/checkpoint.pth.tar'
assert os.path.exists(
filenameCheckpoint
), "Error: resume option was used but checkpoint was not found in folder"
checkpoint = torch.load(filenameCheckpoint)
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
best_acc = checkpoint['best_acc']
print("=> Loaded checkpoint at epoch {})".format(checkpoint['epoch']))
#scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5) # set up scheduler ## scheduler 1
lambda1 = lambda epoch: pow(
(1 - ((epoch - 1) / args.num_epochs)), 0.9) ## scheduler 2
scheduler = lr_scheduler.LambdaLR(optimizer,
lr_lambda=lambda1) ## scheduler 2
if args.visualize and args.steps_plot > 0:
board = Dashboard(args.port)
for epoch in range(start_epoch, args.num_epochs + 1):
print("----- TRAINING - EPOCH", epoch, "-----")
scheduler.step(epoch) ## scheduler 2
epoch_loss = []
time_train = []
doIouTrain = args.iouTrain
doIouVal = args.iouVal
if (doIouTrain):
iouEvalTrain = iouEval(NUM_CLASSES)
usedLr = 0
for param_group in optimizer.param_groups:
print("LEARNING RATE: ", param_group['lr'])
usedLr = float(param_group['lr'])
model.train()
for step, (images, labels) in enumerate(loader):
print(step)
start_time = time.time()
#print (labels.size())
#print (np.unique(labels.numpy()))
#print("labels: ", np.unique(labels[0].numpy()))
#labels = torch.ones(4, 1, 512, 1024).long()
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(images)
targets = Variable(labels)
outputs = model(inputs, only_encode=enc)
#print("targets", np.unique(targets[:, 0].cpu().data.numpy()))
optimizer.zero_grad()
loss = criterion(outputs, targets[:, 0])
loss.backward()
optimizer.step()
epoch_loss.append(loss.data[0])
time_train.append(time.time() - start_time)
print(loss.data[0])
if (doIouTrain):
#start_time_iou = time.time()
iouEvalTrain.addBatch(
outputs.max(1)[1].unsqueeze(1).data, targets.data)
#print ("Time to add confusion matrix: ", time.time() - start_time_iou)
#print(outputs.size())
if args.visualize and args.steps_plot > 0 and step % args.steps_plot == 0:
start_time_plot = time.time()
image = inputs[0].cpu().data
#image[0] = image[0] * .229 + .485
#image[1] = image[1] * .224 + .456
#image[2] = image[2] * .225 + .406
#print("output", np.unique(outputs[0].cpu().max(0)[1].data.numpy()))
board.image(image, 'input (epoch: {epoch}, step: {step})')
if isinstance(outputs, list): #merge gpu tensors
board.image(
color_transform(
outputs[0][0].cpu().max(0)[1].data.unsqueeze(0)),
'output (epoch: {epoch}, step: {step})')
else:
board.image(
color_transform(
outputs[0].cpu().max(0)[1].data.unsqueeze(0)),
'output (epoch: {epoch}, step: {step})')
board.image(color_transform(targets[0].cpu().data),
'target (epoch: {epoch}, step: {step})')
print("Time to paint images: ", time.time() - start_time_plot)
if args.steps_loss > 0 and step % args.steps_loss == 0:
average = sum(epoch_loss) / len(epoch_loss)
print(
'loss: {average:0.4} (epoch: {epoch}, step: {step})',
"// Avg time/img: %.4f s" %
(sum(time_train) / len(time_train) / args.batch_size))
average_epoch_loss_train = sum(epoch_loss) / len(epoch_loss)
iouTrain = 0
if (doIouTrain):
iouTrain, iou_classes = iouEvalTrain.getIoU()
iouStr = getColorEntry(iouTrain) + '{:0.2f}'.format(
iouTrain * 100) + '\033[0m'
print("EPOCH IoU on TRAIN set: ", iouStr, "%")
#Validate on 500 val images after each epoch of training
print("----- VALIDATING - EPOCH", epoch, "-----")
model.eval()
epoch_loss_val = []
time_val = []
if (doIouVal):
iouEvalVal = iouEval(NUM_CLASSES)
for step, (images, labels) in enumerate(loader_val):
start_time = time.time()
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(
images, volatile=True
) #volatile flag makes it free backward or outputs for eval
targets = Variable(labels, volatile=True)
outputs = model(inputs, only_encode=enc)
loss = criterion(outputs, targets[:, 0])
epoch_loss_val.append(loss.data[0])
time_val.append(time.time() - start_time)
#Add batch to calculate TP, FP and FN for iou estimation
if (doIouVal):
#start_time_iou = time.time()
iouEvalVal.addBatch(
outputs.max(1)[1].unsqueeze(1).data, targets.data)
#print ("Time to add confusion matrix: ", time.time() - start_time_iou)
if args.visualize and args.steps_plot > 0 and step % args.steps_plot == 0:
start_time_plot = time.time()
image = inputs[0].cpu().data
board.image(image, 'VAL input (epoch: {epoch}, step: {step})')
if isinstance(outputs, list): #merge gpu tensors
board.image(
color_transform(
outputs[0][0].cpu().max(0)[1].data.unsqueeze(0)),
'VAL output (epoch: {epoch}, step: {step})')
else:
board.image(
color_transform(
outputs[0].cpu().max(0)[1].data.unsqueeze(0)),
'VAL output (epoch: {epoch}, step: {step})')
board.image(color_transform(targets[0].cpu().data),
'VAL target (epoch: {epoch}, step: {step})')
print("Time to paint images: ", time.time() - start_time_plot)
if args.steps_loss > 0 and step % args.steps_loss == 0:
average = sum(epoch_loss_val) / len(epoch_loss_val)
print(
'VAL loss: {average:0.4} (epoch: {epoch}, step: {step})',
"// Avg time/img: %.4f s" %
(sum(time_val) / len(time_val) / args.batch_size))
average_epoch_loss_val = sum(epoch_loss_val) / len(epoch_loss_val)
#scheduler.step(average_epoch_loss_val, epoch) ## scheduler 1 # update lr if needed
iouVal = 0
if (doIouVal):
iouVal, iou_classes = iouEvalVal.getIoU()
iouStr = getColorEntry(iouVal) + '{:0.2f}'.format(
iouVal * 100) + '\033[0m'
print("EPOCH IoU on VAL set: ", iouStr, "%")
# remember best valIoU and save checkpoint
if iouVal == 0:
current_acc = average_epoch_loss_val
else:
current_acc = iouVal
is_best = current_acc > best_acc
best_acc = max(current_acc, best_acc)
if enc:
filenameCheckpoint = savedir + '/checkpoint_enc.pth.tar'
filenameBest = savedir + '/model_best_enc.pth.tar'
else:
filenameCheckpoint = savedir + '/checkpoint.pth.tar'
filenameBest = savedir + '/model_best.pth.tar'
save_checkpoint(
{
'epoch': epoch + 1,
'arch': str(model),
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
}, is_best, filenameCheckpoint, filenameBest)
#SAVE MODEL AFTER EPOCH
if (enc):
filename = savedir + '/model_encoder-' + '{epoch:03}.pth'
filenamebest = savedir + '/model_encoder_best.pth'
else:
filename = savedir + '/model-{epoch:03}.pth'
filenamebest = savedir + '/model_best.pth'
if args.epochs_save > 0 and step > 0 and step % args.epochs_save == 0:
torch.save(model.state_dict(), filename)
print('save: {filename} (epoch: {epoch})')
if (is_best):
torch.save(model.state_dict(), filenamebest)
print('save: {filenamebest} (epoch: {epoch})')
if (not enc):
with open(savedir + "/best.txt", "w") as myfile:
myfile.write("Best epoch is %d, with Val-IoU= %.4f" %
(epoch, iouVal))
else:
with open(savedir + "/best_encoder.txt", "w") as myfile:
myfile.write("Best epoch is %d, with Val-IoU= %.4f" %
(epoch, iouVal))
#SAVE TO FILE A ROW WITH THE EPOCH RESULT (train loss, val loss, train IoU, val IoU)
#Epoch Train-loss Test-loss Train-IoU Test-IoU learningRate
with open(automated_log_path, "a") as myfile:
myfile.write("\n%d\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.8f" %
(epoch, average_epoch_loss_train,
average_epoch_loss_val, iouTrain, iouVal, usedLr))
return (model) #return model (convenience for encoder-decoder training)
def train(args, model):
model.train()
weight = torch.ones(NUM_CLASSES)
#weight[0] = 0
loader = DataLoader(CityScapes(args.datadir, input_transform, target_transform),
num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True)
val_loader = DataLoader(CityScapes_validation(args.datadir, input_transform, target_transform),
num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True)
if args.cuda:
criterion = CrossEntropyLoss2d(weight.cuda())
#criterion=torch.nn.BCEWithLogitsLoss()
else:
criterion = CrossEntropyLoss2d(weight)
optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5)
'''
if args.model.startswith('FCN'):
optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5)
if args.model.startswith('PSP'):
optimizer=SGD(filter(lambda p: p.requires_grad, model.parameters()), 1e-2,0.9,1e-4)
#optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4)
if args.model.startswith('Seg'):
optimizer = SGD(filter(lambda p: p.requires_grad, model.parameters()), 1e-3, .9)
'''
print("Total steps:",len(loader))
best_loss=100
best_val_loss=100
best_acc=0
best_val_acc=0
for epoch in range(1, args.num_epochs+1):
epoch_loss = []
iteration=1
train_acc=[]
for step, (images, labels) in enumerate(loader):
print("Iter:"+str(iteration))
iteration=iteration+1
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(images)
targets = Variable(labels)
outputs=model(inputs)
optimizer.zero_grad()
loss = criterion(outputs, targets[:, 0,:,:])#Bx1xHxW
loss.backward()
optimizer.step()
print(loss.item())
epoch_loss.append(loss.item())
acc_vec=[]
for b in range(outputs.size()[0]):
acc_vec.append(pixel_accuracy(torch.max(outputs[b,:,:,:],0)[1],targets[b,0,:,:]))
acc=sum(acc_vec)/len(acc_vec)
print("train_acc: "+str(acc))
train_acc.append(acc)
if args.steps_loss > 0 and step>0 and step % args.steps_loss == 0:
average = sum(epoch_loss) / len(epoch_loss)
average_acc=sum(train_acc) / len(train_acc)
epoch_loss=[]
train_acc=[]
if best_loss>average:
best_loss=average
#torch.save(model.state_dict(), "model_linknet34.pth")
#print("Model saved!")
if best_acc<average_acc:
best_acc=average_acc
print("loss: "+str(average)+" epoch: "+str(epoch)+", step: "+str(step))
print("best loss: "+str(best_loss)+" epoch: "+str(epoch)+", step: "+str(step))
print("train acc: "+str(average_acc))
print("best train acc: "+str(best_acc))
f=open("train_loss.txt","a")
f.write(str(epoch)+" "+str(step)+" "+str(average)+" "+str(best_loss)+" "+str(average_acc)+" "+str(best_acc)+"\n")
f.close()
print("Best loss: "+str(best_loss))
print("Best val loss: "+str(best_val_loss))
print("best train acc: "+str(best_acc))
print("Best val acc: "+str(best_val_acc))
epoch_loss = []
val_acc=[]
iteration=1
for step, (images, labels) in enumerate(val_loader):
print("Val Iter:"+str(iteration))
iteration=iteration+1
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(images)
targets = Variable(labels)
outputs=model(inputs)
loss = criterion(outputs, targets[:, 0,:,:])
print(loss.item())
epoch_loss.append(loss.item())
val_acc_vec=[]
for b in range(outputs.size()[0]):
val_acc_vec.append(pixel_accuracy(torch.max(outputs[b,:,:,:],0)[1],targets[b,0,:,:]))
acc=sum(val_acc_vec)/len(val_acc_vec)
val_acc.append(acc)
if args.steps_loss > 0 and step>0 and step % args.steps_loss == 0:
average = sum(epoch_loss) / len(epoch_loss)
average_acc=sum(val_acc) / len(val_acc)
epoch_loss=[]
val_acc=[]
if best_val_loss>average:
best_val_loss=average
torch.save(model.state_dict(), "model_linknet34.pth")
print("Model saved!")
if best_val_acc<average_acc:
best_val_acc=average_acc
print("val loss: "+str(average)+" epoch: "+str(epoch)+", step: "+str(step))
print("best val loss: "+str(best_val_loss)+" epoch: "+str(epoch)+", step: "+str(step))
print("val acc: "+str(average_acc))
print("best val acc: "+str(best_acc))
f1=open("val_loss.txt","a")
f1.write(str(epoch)+" "+str(step)+" "+str(average)+" "+str(best_val_loss)+" "+ str(average_acc)+" "+str(best_val_acc)+"\n")
f1.close()
print("Best val loss: "+str(best_val_loss))
print("Best val acc: "+str(best_val_acc))
loader = data.DataLoader(VOCDataSet(
"/home/zeng/data/datasets/segmentation_Dataset",
img_transform=input_transform,
label_transform=target_transform),
batch_size=12,
shuffle=True,
pin_memory=True)
res101 = resnet101(pretrained=True).cuda()
seg = Seg().cuda()
weight = torch.ones(22)
weight[21] = 0
criterion = CrossEntropyLoss2d(weight.cuda())
optimizer_seg = torch.optim.Adam(seg.parameters(), lr=1e-3)
optimizer_feat = torch.optim.Adam(res101.parameters(), lr=1e-4)
for t in range(10):
for i, (img, label) in enumerate(loader):
img = img.cuda()
label = label[0].cuda()
label = Variable(label)
input = Variable(img)
feats = res101(input)
output = seg(feats)
seg.zero_grad()
feature.load_state_dict(torch.load(pretrained_feature_file))
deconv = Deconv(opt.i)
deconv.cuda()
if resume_ep >= 0:
feature_param_file = glob.glob('%s/feature-epoch-%d*.pth'%(check_dir, resume_ep))
deconv_param_file = glob.glob('%s/deconv-epoch-%d*.pth'%(check_dir, resume_ep))
feature.load_state_dict(torch.load(feature_param_file[0]))
deconv.load_state_dict(torch.load(deconv_param_file[0]))
train_loader = torch.utils.data.DataLoader(
MyBoxPixData(train_dir, transform=True, crop=True, hflip=True, vflip=False, source=opt.q),
batch_size=bsize, shuffle=True, num_workers=4, pin_memory=True)
criterion = CrossEntropyLoss2d(weight=torch.FloatTensor(label_weight))
criterion.cuda()
optimizer_deconv = torch.optim.Adam(deconv.parameters(), lr=1e-3)
optimizer_feature = torch.optim.Adam(feature.parameters(), lr=1e-4)
for it in range(resume_ep+1, iter_num):
for ib, (data, lbl) in enumerate(train_loader):
inputs = Variable(data).cuda()
lbl = Variable(lbl.long()).cuda()
feats = feature(inputs)
msk = deconv(feats)
msk = functional.upsample(msk, scale_factor=8)
loss = criterion(msk, lbl)
def train(args, model):
model.train()
weight = torch.ones(22)
weight[0] = 0
loader = DataLoader(VOC12(args.datadir, input_transform, target_transform),
num_workers=args.num_workers,
batch_size=args.batch_size,
shuffle=True)
if args.cuda:
criterion = CrossEntropyLoss2d(weight.cuda())
else:
criterion = CrossEntropyLoss2d(weight)
optimizer = Adam(model.parameters())
if args.model.startswith('FCN'):
optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5)
if args.model.startswith('PSP'):
optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4)
if args.model.startswith('Seg'):
optimizer = SGD(model.parameters(), 1e-3, .9)
if args.steps_plot > 0:
board = Dashboard(args.port)
for epoch in range(1, args.num_epochs + 1):
epoch_loss = []
for step, (images, labels) in enumerate(loader):
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(images)
targets = Variable(labels)
outputs = model(inputs)
optimizer.zero_grad()
loss = criterion(outputs, targets[:, 0])
loss.backward()
optimizer.step()
epoch_loss.append(loss.data[0])
if args.steps_plot > 0 and step % args.steps_plot == 0:
image = inputs[0].cpu().data
image[0] = image[0] * .229 + .485
image[1] = image[1] * .224 + .456
image[2] = image[2] * .225 + .406
board.image(image, f('input (epoch: {epoch}, step: {step})'))
board.image(color_transform(outputs[0].cpu().max(0)[1].data),
f('output (epoch: {epoch}, step: {step})'))
board.image(color_transform(targets[0].cpu().data),
f('target (epoch: {epoch}, step: {step})'))
if args.steps_loss > 0 and step % args.steps_loss == 0:
average = sum(epoch_loss) / len(epoch_loss)
print(f('loss: {average} (epoch: {epoch}, step: {step})'))
if args.steps_save > 0 and step % args.steps_save == 0:
filename = f('{args.model}-{epoch:03}-{step:04}.pth')
torch.save(model.state_dict(), filename)
print(f('save: {filename} (epoch: {epoch}, step: {step})'))
def train(args, model):
model.train()
weight = torch.ones(22)
weight[0] = 0
loader = DataLoader(VOC12(args.datadir, input_transform, target_transform),
num_workers=args.num_workers, batch_size=args.batch_size, shuffle=True)
if args.cuda:
criterion = CrossEntropyLoss2d(weight.cuda())
else:
criterion = CrossEntropyLoss2d(weight)
#optimizer = Adam(model.parameters())
optimizer=SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)
'''
if args.model.startswith('FCN'):
optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5)
if args.model.startswith('PSP'):
optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4)
if args.model.startswith('Seg'):
optimizer = SGD(model.parameters(), 1e-3, .9)
'''
'''
if args.steps_plot > 0:
board = Dashboard(args.port)
'''
print(len(loader))
for epoch in range(1, args.num_epochs+1):
epoch_loss = []
iteration=1
for step, (images, labels) in enumerate(loader):
print("Iter:"+str(iteration))
iteration=iteration+1
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(images)
targets = Variable(labels)
#change
outputs = model(inputs)
optimizer.zero_grad()
loss = criterion(outputs, targets[:, 0])
loss.backward()
optimizer.step()
epoch_loss.append(loss.data[0])
if args.steps_plot > 0 and step % args.steps_plot == 0:
image = inputs[0].cpu().data
image[0] = image[0] * .229 + .485
image[1] = image[1] * .224 + .456
image[2] = image[2] * .225 + .406
'''
board.image(image,
f'input (epoch: {epoch}, step: {step})')
board.image(color_transform(outputs[0].cpu().max(0)[1].data),
f'output (epoch: {epoch}, step: {step})')
board.image(color_transform(targets[0].cpu().data),
f'target (epoch: {epoch}, step: {step})')
'''
#if args.steps_loss > 0 and step % args.steps_loss == 0:
average = sum(epoch_loss) / len(epoch_loss)
print("loss: "+str(average)+" epoch: "+str(epoch)+", step: "+str(step))
'''
seg_loader = torch.utils.data.DataLoader(MyBoxPixData(seg_train_dir,
transform=True,
crop=True,
hflip=True,
vflip=False,
source=opt.q),
batch_size=bsize,
shuffle=True,
num_workers=4,
pin_memory=True)
criterion_cls = nn.CrossEntropyLoss(weight=torch.FloatTensor(cls_label_weight))
criterion_cls.cuda()
criterion_seg = CrossEntropyLoss2d(weight=torch.FloatTensor(seg_label_weight))
criterion_seg.cuda()
optimizer_classifier = torch.optim.Adam(classifier.parameters(), lr=1e-3)
optimizer_deconv = torch.optim.Adam(deconv.parameters(), lr=1e-3)
optimizer_feature = torch.optim.Adam(feature.parameters(), lr=1e-4)
segIter = iter(seg_loader)
ibs = 0
for it in range(resume_ep + 1, iter_num):
for ib, (data, lbl) in enumerate(cls_loader):
# train with classification data
inputs = Variable(data.float()).cuda()
lbl = Variable(lbl.long()).cuda()
feats = feature(inputs)
output = classifier(feats)
batch_size=1, shuffle=True, num_workers=4, pin_memory=True)
else:
outputRoot = opt.out
loader = torch.utils.data.DataLoader(
MyTestData(dataRoot, transform=True),
batch_size=1, shuffle=True, num_workers=4, pin_memory=True)
"""nets"""
model = RFCN()
if opt.param is None:
vgg16 = torchvision.models.vgg16(pretrained=True)
model.copy_params_from_vgg16(vgg16, copy_fc8=False, init_upscore=True)
else:
model.load_state_dict(torch.load(opt.param))
criterion = CrossEntropyLoss2d()
optimizer = torch.optim.Adam(model.parameters(), 0.0001, betas=(0.5, 0.999))
model = model.cuda()
mr_sal = MR.MR_saliency()
if opt.phase == 'train':
"""train"""
for it in range(iterNum):
epoch_loss = []
for ib, data in enumerate(loader):
# prior map
_img = tensor2image(data[0][0])
pmap = mr_sal.saliency(_img).astype(float) / 255.0
pmap = 1.0 - pmap
pmap = torch.unsqueeze(torch.FloatTensor(pmap), 0)
def train(args, model):
# set model to training mode
model.train()
# prepare criterion
weight = torch.ones(22)
weight[0] = 0
train_loader = DataLoader(VOC12(root=args.datadir,
train=True,
input_transform=input_transform,
target_transform=target_transform),
num_workers=args.num_workers,
batch_size=args.batch_size,
shuffle=False)
val_loader = DataLoader(VOC12(root=args.datadir,
train=False,
input_transform=input_transform,
target_transform=target_transform),
num_workers=args.num_workers,
batch_size=args.batch_size,
shuffle=False)
# use Adam optimizer
optimizer = Adam(model.parameters())
# for loss calculation, we still use CrossEntropyLoss
if args.cuda:
criterion = CrossEntropyLoss2d(weight.cuda())
else:
criterion = CrossEntropyLoss2d(weight)
# start training - epoch values start from 1 to make numbers look 'pretty'
for epoch in range(1, args.num_epochs + 1):
epoch_loss = []
print(len(train_loader))
for step, (images, labels) in enumerate(train_loader):
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(images)
targets = Variable(labels)
outputs = model(inputs)
# refresh gradient before backprop
optimizer.zero_grad()
loss = criterion(outputs, targets[:, 0])
loss.backward()
optimizer.step()
epoch_loss.append(loss.item())
if (args.steps_loss > 0) and (step % args.steps_loss == 0):
# print loss
average = sum(epoch_loss) / len(epoch_loss)
print(f'loss: {average} (epoch {epoch}, step {step})')
# check for early stop
if metrics.early_stop(use_cuda=args.cuda,
output=outputs,
target=targets[:, 0],
metric="iou",
threshold=0.8):
# if we can stop early, save model and exit
print(
"Early stopping score exceeded threshold... saving model and ending training stage"
)
if args.attention:
filename = f'models/AttentionUNet-EarlyStop-{epoch:03}-{step:04}.pth'
else:
filename = f'models/UNet-EarlyStop-{epoch:03}-{step:04}.pth'
torch.save(model.state_dict(), filename)
return
if (args.steps_save > 0) and (step % args.steps_save == 0):
if args.attention:
filename = f'models/AttentionUNet-{epoch:03}-{step:04}.pth'
else:
filename = f'models/UNet-{epoch:03}-{step:04}.pth'
torch.save(model.state_dict(), filename)
print(f'save: {filename} (epoch: {epoch}, step: {step})')
''' printing out pictures doesn't work on VScode :(
_, outputs = torch.max(outputs, dim=1)
#print(outputs.shape)
outputs = outputs.unsqueeze(1)
#print(outputs.shape)
# show images
plt.subplot(211)
plt.imshow(make_grid(images.cpu()).permute(1,2,0).numpy())
plt.axis('off')
plt.title('Images')
# print labels
plt.subplot(212)
outputs = outputs.cpu().numpy()[:, :, :, :, np.newaxis]
color_Label = np.dot(outputs == 0, cmap[0])
for i in range(1, cmap.shape[0]):
color_Label += np.dot(outputs == i, cmap[i])
color_Label = color_Label.swapaxes(1,4)
plt.imshow((make_grid(torch.tensor(color_Label.squeeze())).permute(1,2,0).numpy()).astype('uint8'))
plt.axis('off')
plt.title('Label')
'''
# every epoch, check validation data's accuracy
v_list = []
for _, (v_images, v_labels) in enumerate(val_loader):
if args.cuda:
v_images = v_images.cuda()
v_labels = v_labels.cuda()
v_inputs = Variable(v_images)
v_targets = Variable(v_labels)
v_outputs = model(v_inputs)
# using cross-entropy loss
# v_loss = criterion(v_outputs, v_targets[:, 0])
# v_list.append(v_loss.item())
# for IoU metric
_, v_l = torch.max(v_outputs, dim=1)
v_score = metrics.get_iou(labels=v_l, target=v_targets[:, 0])
v_list.append(v_score)
v_average = sum(v_list) / len(v_list)
print(f'validation average IoU: {v_average} (epoch {epoch})')
if __name__ == '__main__':
models_path = os.path.join('./checkpoints', args.backend) #保存网络参数的路径
os.makedirs(models_path, exist_ok=True)
train_loader, val_loader = get_dataloader()
net, starting_epoch = build_network(args.snapshot, args.backend)
optimizer = optim.Adam(net.parameters(), lr=args.start_lr) #优化器
# optimizer = optim.Adam(net.parameters())
# optimizer = optim.SGD(net.parameters(), lr=1e-3, momentum=0.9,
# weight_decay=1e-4)
scheduler = MultiStepLR(
optimizer, milestones=[int(x) for x in args.milestones.split(',')])
seg_criterion = CrossEntropyLoss2d(weight=None) #损失函数
epoch_losses = []
net.train()
weight_save_path = "checkpoints/densenet/PSPNet_last" #网络参数地址
#加载网络参数
try:
net.load_state_dict(torch.load(weight_save_path))
print("加载成功")
except:
print('加载失败')
#开始训练
Loss_list = []
Accuracy_list = []
for epoch in range(1 + starting_epoch, 1 + starting_epoch + args.epochs):
def train(args, model):
model.train()
weight = torch.ones(NUM_CLASSES)
weight[0] = 0
loader = DataLoader(VOC12(args.datadir, input_transform, target_transform),
num_workers=args.num_workers,
batch_size=args.batch_size,
shuffle=True)
if args.cuda:
criterion = CrossEntropyLoss2d(weight.cuda())
#criterion=torch.nn.BCEWithLogitsLoss()
else:
criterion = CrossEntropyLoss2d(weight)
if args.model.startswith('FCN'):
optimizer = SGD(model.parameters(), 1e-4, .9, 2e-5)
if args.model.startswith('PSP'):
optimizer = SGD(filter(lambda p: p.requires_grad, model.parameters()),
1e-2, 0.9, 1e-4)
#optimizer = SGD(model.parameters(), 1e-2, .9, 1e-4)
if args.model.startswith('Seg'):
optimizer = SGD(filter(lambda p: p.requires_grad, model.parameters()),
1e-3, .9)
print("Total images:", len(loader))
best_loss = 100
f = open("loss.txt", "a")
for epoch in range(1, args.num_epochs + 1):
epoch_loss = []
iteration = 1
for step, (images, labels) in enumerate(loader):
print("Iter:" + str(iteration))
iteration = iteration + 1
if args.cuda:
images = images.cuda()
labels = labels.cuda()
inputs = Variable(images)
targets = Variable(labels)
outputs = model(inputs)
optimizer.zero_grad()
loss = criterion(outputs, targets[:, 0])
loss.backward()
optimizer.step()
print(loss.data[0])
epoch_loss.append(loss.data[0])
if args.steps_loss > 0 and step > 0 and step % args.steps_loss == 0:
average = sum(epoch_loss) / len(epoch_loss)
epoch_loss = []
if best_loss > average:
best_loss = average
torch.save(model.state_dict(),
"model_pspnet_VOC_2012_analysis.pth")
print("Model saved!")
f.write("loss: " + str(average) + " epoch: " + str(epoch) +
", step: " + str(step) + "\n")
f.write("best loss: " + str(best_loss) + " epoch: " +
str(epoch) + ", step: " + str(step) + "\n")
print("loss: " + str(average) + " epoch: " + str(epoch) +
", step: " + str(step))
print("best loss: " + str(best_loss) + " epoch: " +
str(epoch) + ", step: " + str(step))
print("Best loss: " + str(best_loss))