model = SegNetBasic(args.in_channel, args.n_classes) elif args.model == 'UNet': model = UNet(args.in_channel, args.n_classes) else: print('This model doesn\'t exist in the model directory') sys.exit(1) if args.params_path is not None: model.load_state_dict( torch.load(args.params_path, map_location=lambda storage, loc: storage)) """ define DataLoader """ data = PartAffordanceDataset('test.csv', transform=transforms.Compose( [CenterCrop(), ToTensor(), Normalize()])) data_loader = DataLoader(data, batch_size=args.num_images, shuffle=True) for sample in data_loader: model.eval() predict(model, sample, device=args.device) x = sample["image"] x = reverse_normalize(x, mean, std) save_image( x, args.result_path + '/' + 'original_images_with_' + args.model + '.jpg')
import torch from torch.utils.data import DataLoader from torch.nn import functional as F from torchvision.transforms import Compose from model import UNet from dataset import Segmentation, RandomAffine, Pad, RandomFlip, CenterCrop, ToTensor, RandomWarp from torchvision import transforms as T from PIL import Image import numpy as np # Dataset dataset = Segmentation(transform = Compose([ \ Pad(150, mode='symmetric'), \ RandomAffine((0, 90), (30, 30)), \ CenterCrop(512, 512), \ RandomFlip(), \ RandomWarp(), CenterCrop(512, 504), \ ToTensor() ])) # Neural network model = UNet(n_class = 1).cuda() if torch.cuda.is_available() else UNet(n_class = 1) def save_checkpoint(checkpt, filename): torch.save(checkpt,filename) def get_checkpoint(model, optimizer, loss): filename = "unet.pth" map_location = 'cuda:0' if torch.cuda.is_available() else 'cpu'
def main(): args = get_arguments() # configuration CONFIG = Dict(yaml.safe_load(open(args.config))) """ DataLoader """ test_data = PartAffordanceDataset(CONFIG.test_data, config=CONFIG, transform=transforms.Compose([ CenterCrop(CONFIG), ToTensor(), Normalize() ])) test_loader = DataLoader(test_data, batch_size=4, shuffle=True, num_workers=1) test_iter = iter(test_loader) model = models.vgg16_bn(pretrained=False) model.classifier[6] = nn.Linear(in_features=4096, out_features=7, bias=True) model.load_state_dict( torch.load('./result/best_accuracy_model.prm', map_location=lambda storage, loc: storage)) model.to(args.device) while True: sample = test_iter.next() image, label = sample['image'], sample['label'] # show images show_img(torchvision.utils.make_grid(image)) # print labels print('True labels') print(label) with torch.no_grad(): image = image.to(args.device) label = label.to(args.device) h = model(image) h = torch.sigmoid(h) h[h > 0.5] = 1 h[h <= 0.5] = 0 total_num = 7 * len(label) acc_num = torch.sum(h == label) accuracy = float(acc_num) / total_num print('\nPredicted labels') print(h) print('\naccuracy\t{:.3f}'.format(accuracy)) print('\nIf you want to look at more images, input \"c\"') s = input() if s == 'c': continue else: break
def main(): """ DataLoader """ train_data = PartAffordanceDataset( 'train.csv', transform=transforms.Compose([CenterCrop(), ToTensor(), Normalize()])) test_data = PartAffordanceDataset( 'test.csv', transform=transforms.Compose([CenterCrop(), ToTensor(), Normalize()])) train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, num_workers=4) test_loader = DataLoader(test_data, batch_size=args.batch_size, shuffle=False, num_workers=4) if args.model == 'FCN8s': model = FCN8s(args.in_channel, args.n_classes) elif args.model == 'SegNetBasic': model = SegNetBasic(args.in_channel, args.n_classes) elif args.model == 'UNet': model = UNet(args.in_channel, args.n_classes) else: print('This model doesn\'t exist in the model directory') sys.exit(1) model.apply(init_weight) """ training """ if args.writer: writer = SummaryWriter(args.result_path) if args.class_weight: criterion = nn.CrossEntropyLoss(weight=class_weight.to(args.device)) else: criterion = nn.CrossEntropyLoss() model.to(args.device) optimizer = optim.Adam(model.parameters(), lr=args.learning_rate) train_losses = [] val_iou = [] mean_iou = [] best_mean_iou = 0.0 for epoch in range(args.max_epoch): model.train() running_loss = 0.0 for i, sample in tqdm.tqdm(enumerate(train_loader), total=len(train_loader)): optimizer.zero_grad() x, y = sample['image'], sample['class'] x = x.to(args.device) y = y.to(args.device) h = model(x) loss = criterion(h, y) loss.backward() optimizer.step() running_loss += loss.item() train_losses.append(running_loss / i) val_iou.append( eval_model(model, test_loader, args.device).to('cpu').float()) mean_iou.append(val_iou[-1].mean().item()) if best_mean_iou < mean_iou[-1]: best_mean_iou = mean_iou[-1] torch.save(model.state_dict(), args.result_path + '/best_mean_iou_model.prm') if writer is not None: writer.add_scalar("train_loss", train_losses[-1], epoch) writer.add_scalar("mean_IoU", mean_iou[-1], epoch) writer.add_scalars( "class_IoU", { 'iou of class 0': val_iou[-1][0], 'iou of class 1': val_iou[-1][1], 'iou of class 2': val_iou[-1][2], 'iou of class 3': val_iou[-1][3], 'iou of class 4': val_iou[-1][4], 'iou of class 5': val_iou[-1][5], 'iou of class 6': val_iou[-1][6], 'iou of class 7': val_iou[-1][7] }, epoch) print('epoch: {}\tloss: {:.5f}\tmean IOU: {:.3f}'.format( epoch, train_losses[-1], mean_iou[-1])) torch.save(model.state_dict(), args.result_path + "/final_model.prm")
import torch from torchvision import transforms as T from PIL import Image from dataset import Segmentation, RandomFlip, Pad, RandomAffine, CenterCrop, ToTensor, RandomWarp from model import UNet from dd import deform_grid t = T.Compose([ #RandomFlip(), Pad(150, mode='symmetric'), #RandomAffine((0, 90), (31, 31)), #RandomWarp(), CenterCrop(572, 388), ToTensor() ]) transform = T.Compose([ \ Pad(150, mode='symmetric'), \ RandomAffine((0, 90), (31, 31)), \ RandomFlip(), \ RandomWarp(), CenterCrop(572, 388), \ ToTensor() ]) dataset = Segmentation(transform=t) original = Segmentation(transform=ToTensor()) pil = T.ToPILImage() def seg(sample):
from torch.utils.data import DataLoader from torch.nn import functional as F from torchvision.transforms import Compose from model import UNet from dataset import Segmentation, RandomAffine, Pad, RandomFlip, CenterCrop, ToTensor, RandomWarp from torchvision import transforms as T from PIL import Image import numpy as np # Dataset dataset = Segmentation(transform = Compose([ \ Pad(150, mode='symmetric'), \ RandomAffine((0, 90), (30, 30)), \ CenterCrop(512, 512), \ RandomFlip(), \ RandomWarp(), CenterCrop(512, 504), \ ToTensor() ])) # Neural network model = UNet(n_class=1).cuda() if torch.cuda.is_available() else UNet( n_class=1) def save_checkpoint(checkpt, filename): torch.save(checkpt, filename) def get_checkpoint(model, optimizer, loss):
def main(): args = get_arguments() # configuration CONFIG = Dict(yaml.safe_load(open(args.config))) # writer if CONFIG.writer_flag: writer = SummaryWriter(CONFIG.result_path) else: writer = None """ DataLoader """ labeled_train_data = PartAffordanceDataset(CONFIG.labeled_data, config=CONFIG, transform=transforms.Compose([ CenterCrop(CONFIG), ToTensor(), Normalize() ])) if CONFIG.train_mode == 'semi': unlabeled_train_data = PartAffordanceDatasetWithoutLabel(CONFIG.unlabeled_data, config=CONFIG, transform=transforms.Compose([ CenterCrop(CONFIG), ToTensor(), Normalize() ])) else: unlabeled_train_data = None test_data = PartAffordanceDataset(CONFIG.test_data, config=CONFIG, transform=transforms.Compose([ CenterCrop(CONFIG), ToTensor(), Normalize() ])) train_loader_with_label = DataLoader(labeled_train_data, batch_size=CONFIG.batch_size, shuffle=True, num_workers=CONFIG.num_workers) if unlabeled_train_data is not None: train_loader_without_label = DataLoader(unlabeled_train_data, batch_size=CONFIG.batch_size, shuffle=True, num_workers=CONFIG.num_workers) test_loader = DataLoader(test_data, batch_size=CONFIG.batch_size, shuffle=False, num_workers=CONFIG.num_workers) """ model """ if CONFIG.model == 'FCN8s': model = FCN8s(CONFIG.in_channel, CONFIG.n_classes) elif CONFIG.model == 'SegNetBasic': model = SegNetBasic(CONFIG.in_channel, CONFIG.n_classes) elif CONFIG.model == 'UNet': model = UNet(CONFIG.in_channel, CONFIG.n_classes) else: print('This model doesn\'t exist in the model directory') sys.exit(1) if CONFIG.train_mode == 'full': model.apply(init_weights) model.to(args.device) elif CONFIG.train_mode == 'semi': if CONFIG.pretrain_model is not None: torch.load(CONFIG.pretrain_model, map_location=lambda storage, loc: storage) else: model.apply(init_weights) model.to(args.device) if CONFIG.model_d == 'FCDiscriminator': model_d = FCDiscriminator(CONFIG) else: model_d = Discriminator(CONFIG) model_d.apply(init_weights) model_d.to(args.device) else: print('This training mode doesn\'t exist.') sys.exit(1) """ class weight after center crop. See dataset.py """ if CONFIG.class_weight_flag: class_num = torch.tensor([2078085712, 34078992, 15921090, 12433420, 38473752, 6773528, 9273826, 20102080]) total = class_num.sum().item() frequency = class_num.float() / total median = torch.median(frequency) class_weight = median / frequency class_weight = class_weight.to(args.device) else: class_weight = None """ supplementary constant for discriminator """ if CONFIG.noisy_label_flag: if one_label_smooth: real = torch.full((CONFIG.batch_size, CONFIG.height, CONFIG.width), CONFIG.real_label).to(args.device) fake = torch.zeros(CONFIG.batch_size, CONFIG.height, CONFIG.width).to(args.device) else: real = torch.full((CONFIG.batch_size, CONFIG.height, CONFIG.width), CONFIG.real_label).to(args.device) fake = torch.full((CONFIG.batch_size, CONFIG.height, CONFIG.width), CONFIG.fake_label).to(args.device) else: real = torch.ones(CONFIG.batch_size, CONFIG.height, CONFIG.width).to(args.device) fake = torch.zeros(CONFIG.batch_size, CONFIG.height, CONFIG.width).to(args.device) """ optimizer, criterion """ optimizer = optim.Adam(model.parameters(), lr=CONFIG.learning_rate) criterion = nn.CrossEntropyLoss(weight=class_weight, ignore_index=255) if CONFIG.train_mode == 'semi': criterion_bce = nn.BCELoss() # discriminator includes sigmoid layer optimizer_d = optim.Adam(model_d.parameters(), lr=CONFIG.learning_rate) losses_full = [] losses_semi = [] losses_d = [] val_iou = [] mean_iou = [] mean_iou_without_bg = [] best_mean_iou = 0.0 for epoch in tqdm.tqdm(range(CONFIG.max_epoch)): if CONFIG.poly_lr_decay: poly_lr_scheduler(optimizer, CONFIG.learning_rate, epoch, max_iter=CONFIG.max_epoch, power=CONFIG.poly_power) if CONFIG.train_mode == 'semi': poly_lr_scheduler(optimizer_d, CONFIG.learning_rate_d, epoch, max_iter=CONFIG.max_epoch, power=CONFIG.poly_power) epoch_loss_full = 0.0 epoch_loss_d = 0.0 epoch_loss_semi = 0.0 # only supervised learning if CONFIG.train_mode == 'full': for i, sample in enumerate(train_loader_with_label): loss_full = full_train(model, sample, criterion, optimizer, CONFIG, args.device) epoch_loss_full += loss_full losses_full.append(epoch_loss_full / i) losses_d.append(0.0) losses_semi.append(0.0) # semi-supervised learning elif CONFIG.train_mode == 'semi': # first, adveresarial learning if epoch < CONFIG.adv_epoch: for i, sample in enumerate(train_loader_with_label): loss_full, loss_d = adv_train( model, model_d, sample, criterion, criterion_bce, optimizer, optimizer_d, real, fake, CONFIG, args.device) epoch_loss_full += loss_full epoch_loss_d += loss_d losses_full.append(epoch_loss_full / i) # mean loss over all samples losses_d.append(epoch_loss_d / i) losses_semi.append(0.0) # semi-supervised learning else: cnt_full = 0 cnt_semi = 0 for (sample1, sample2) in zip_longest(train_loader_with_label, train_loader_without_label): if sample1 is not None: loss_full, loss_d = adv_train( model, model_d, sample1, criterion, criterion_bce, optimizer, optimizer_d, real, fake, CONFIG, args.device) epoch_loss_full += loss_full epoch_loss_d += loss_d cnt_full += 1 if sample2 is not None: loss_semi = semi_train( model, model_d, sample2, criterion, criterion_bce, optimizer, optimizer_d, real, fake, CONFIG, args.device) epoch_loss_semi += loss_semi cnt_semi += 1 losses_full.append(epoch_loss_full / cnt_full) # mean loss over all samples losses_d.append(epoch_loss_d / cnt_full) losses_semi.append(epoch_loss_semi / cnt_semi) else: print('This train mode can\'t be used. Choose full or semi') sys.exit(1) # validation val_iou.append(eval_model(model, test_loader, CONFIG, args.device)) mean_iou.append(val_iou[-1].mean().item()) mean_iou_without_bg.append(val_iou[-1][1:].mean().item()) if best_mean_iou < mean_iou[-1]: best_mean_iou = mean_iou[-1] torch.save(model.state_dict(), CONFIG.result_path + '/best_mean_iou_model.prm') if CONFIG.train_mode == 'semi': torch.save(model_d.state_dict(), CONFIG.result_path + '/best_mean_iou_model_d.prm') if epoch%50 == 0 and epoch != 0: torch.save(model.state_dict(), CONFIG.result_path + '/epoch_{}_model.prm'.format(epoch)) if CONFIG.train_mode == 'semi': torch.save(model_d.state_dict(), CONFIG.result_path + '/epoch_{}_model_d.prm'.format(epoch)) if writer is not None: writer.add_scalar("loss_full", losses_full[-1], epoch) writer.add_scalar("loss_d", losses_d[-1], epoch) writer.add_scalar("loss_semi", losses_semi[-1], epoch) writer.add_scalar("mean_iou", mean_iou[-1], epoch) writer.add_scalar("mean_iou_without_background", mean_iou_without_bg[-1], epoch) writer.add_scalars("class_IoU", {'iou of class 0': val_iou[-1][0], 'iou of class 1': val_iou[-1][1], 'iou of class 2': val_iou[-1][2], 'iou of class 3': val_iou[-1][3], 'iou of class 4': val_iou[-1][4], 'iou of class 5': val_iou[-1][5], 'iou of class 6': val_iou[-1][6], 'iou of class 7': val_iou[-1][7]}, epoch) print('epoch: {}\tloss_full: {:.5f}\tloss_d: {:.5f}\tloss_semi: {:.5f}\tmean IOU: {:.3f}\tmean IOU w/ bg: {:.3f}' .format(epoch, losses_full[-1], losses_d[-1], losses_semi[-1], mean_iou[-1], mean_iou_without_bg[-1])) torch.save(model.state_dict(), CONFIG.result_path + '/final_model.prm') if CONFIG.train_mode == 'semi': torch.save(model_d.state_dict(), CONFIG.result_path + '/final_model_d.prm')
def main(): args = get_arguments() # configuration CONFIG = Dict(yaml.safe_load(open(args.config))) # writer if CONFIG.writer_flag: writer = SummaryWriter(CONFIG.result_path) else: writer = None """ DataLoader """ train_data = PartAffordanceDataset(CONFIG.labeled_data, config=CONFIG, transform=transforms.Compose([ CenterCrop(CONFIG), ToTensor(), Normalize() ])) test_data = PartAffordanceDataset(CONFIG.test_data, config=CONFIG, transform=transforms.Compose([ CenterCrop(CONFIG), ToTensor(), Normalize() ])) train_loader = DataLoader(train_data, batch_size=CONFIG.batch_size, shuffle=True, num_workers=CONFIG.num_workers) test_loader = DataLoader(test_data, batch_size=CONFIG.batch_size, shuffle=False, num_workers=CONFIG.num_workers) model = models.vgg16_bn(pretrained=True) for param in model.features.parameters(): param.requires_grad = False model.classifier[6] = nn.Linear(in_features=4096, out_features=7, bias=True) model.to(args.device) """ optimizer, criterion """ optimizer = optim.Adam(model.classifier.parameters(), lr=CONFIG.learning_rate) criterion = nn.BCEWithLogitsLoss() losses_train = [] losses_val = [] class_accuracy_val = [] accuracy_val = [] best_accuracy = 0.0 for epoch in tqdm.tqdm(range(CONFIG.max_epoch)): poly_lr_scheduler(optimizer, CONFIG.learning_rate, epoch, max_iter=CONFIG.max_epoch, power=CONFIG.poly_power) epoch_loss = 0.0 for sample in train_loader: loss_train = full_train(model, sample, criterion, optimizer, args.device) epoch_loss += loss_train losses_train.append(epoch_loss / len(train_loader)) # validation loss_val, class_accuracy, accuracy = eval_model( model, test_loader, criterion, CONFIG, args.device) losses_val.append(loss_val) class_accuracy_val.append(class_accuracy) accuracy_val.append(accuracy) if best_accuracy < accuracy_val[-1]: best_accuracy = accuracy_val[-1] torch.save(model.state_dict(), CONFIG.result_path + '/best_accuracy_model.prm') if epoch % 50 == 0 and epoch != 0: torch.save( model.state_dict(), CONFIG.result_path + '/epoch_{}_model.prm'.format(epoch)) if writer is not None: writer.add_scalars("loss", { 'loss_train': losses_train[-1], 'loss_val': losses_val[-1] }, epoch) writer.add_scalar("accuracy", accuracy_val[-1], epoch) writer.add_scalars( "class_accuracy", { 'accuracy of class 1': class_accuracy_val[-1][0], 'accuracy of class 2': class_accuracy_val[-1][1], 'accuracy of class 3': class_accuracy_val[-1][2], 'accuracy of class 4': class_accuracy_val[-1][3], 'accuracy of class 5': class_accuracy_val[-1][4], 'accuracy of class 6': class_accuracy_val[-1][5], 'accuracy of class 7': class_accuracy_val[-1][6], }, epoch) print( 'epoch: {}\tloss_train: {:.5f}\tloss_val: {:.5f}\taccuracy: {:.5f}' .format(epoch, losses_train[-1], losses_val[-1], accuracy_val[-1])) torch.save(model.state_dict(), CONFIG.result_path + '/final_model.prm')