def train(data_size='all'): device = torch.device('cuda') if torch.cuda.is_available else torch.device( 'cpu') # Image input model = SegNet(opt, 3) model = model.to(device) model.train() criterion = torch.nn.CrossEntropyLoss() criterion_d = DiscriminativeLoss() optimizer = SGD(model.parameters(), lr=opt.lr, momentum=opt.momentum) if data_size == 'all': dataloader = get_dataloader(opt.paths, opt, device) model = batch_step(opt, optimizer, model, dataloader, criterion, criterion_d, device) torch.save(model.state_dict(), 'model_all.pth') else: im = Image.open(opt.img_path) im = np.array(im, dtype=np.float32) / 255 image = np.transpose(im, (2, 0, 1)) data = torch.from_numpy(image).unsqueeze(0) data = Variable(data).to(device) labels = segmentation.slic(im, compactness=opt.compactness, n_segments=opt.num_superpixels) labels = labels.reshape(-1) label_nums = np.unique(labels) label_indices = [ np.where(labels == label_nums[i])[0] for i in range(len(label_nums)) ] model = one_step(opt, optimizer, model, data, label_indices, criterion, criterion_d, device) torch.save(model.state_dict(), 'model_single.pth')
writer.scalar_summary('train_loss', loss, batches_done) # Determine approximate time left for epoch epoch_batches_left = len(train_dataloader) - (index + 1) time_left = datetime.timedelta(seconds=epoch_batches_left * (time.time() - t_start) / (index + 1)) print('epoch: {}\tbatches: {}\tloss: {:.8f}\tremaining time: {}'. format(epoch, batches_done, loss, time_left)) writer.scalar_summary('train_loss_epoch', epoch_loss.avg, epoch + 1) logger.info('{} epoch loss: {}'.format(epoch + 1, epoch_loss.avg)) is_better = epoch_loss.avg < prev_loss if is_better: prev_loss = epoch_loss.avg torch.save(model.state_dict(), f"checkpoints/best_ckpt_%d.pth" % epoch) else: torch.save(model.state_dict(), f"checkpoints/segnet_ckpt_%d.pth" % epoch) if epoch % args.eval_interval == 0: val_loss, score, class_iou = validate(model=model, val_path=val_path, img_path=val_img_path, mask_path=val_mask_path, batch_size=8) writer.scalar_summary('val_loss', val_loss, epoch + 1) logger.info('epoch {} val loss: {}'.format(epoch + 1, val_loss)) for k, v in score.items(): print(k, v)
def train_autoencoder(epoch_plus): writer = SummaryWriter(log_dir='./runs_autoencoder_2') num_epochs = 400 - epoch_plus lr = 0.001 bta1 = 0.9 bta2 = 0.999 weight_decay = 0.001 # model = autoencoder(nchannels=3, width=172, height=600) model = SegNet(3) if ngpu > 1: model = nn.DataParallel(model) if use_gpu: model = model.to(device, non_blocking=True) if epoch_plus > 0: model.load_state_dict( torch.load('./autoencoder_models_2/autoencoder_{}.pth'.format( epoch_plus))) criterion = nn.MSELoss(reduction='sum') optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(bta1, bta2), weight_decay=weight_decay) for epoch in range(num_epochs): degree = randint(-180, 180) transforms = torchvision.transforms.Compose([ torchvision.transforms.CenterCrop((172, 200)), torchvision.transforms.Resize((172, 200)), torchvision.transforms.RandomRotation((degree, degree)), torchvision.transforms.ToTensor() ]) dataloader = get_dataloader(data_dir, train=True, transform=transforms, batch_size=batch_size) model.train() epoch_losses = AverageMeter() with tqdm(total=(1000 - 1000 % batch_size)) as _tqdm: _tqdm.set_description('epoch: {}/{}'.format( epoch + 1 + epoch_plus, num_epochs + epoch_plus)) for data in dataloader: gt, text = data if use_gpu: gt, text = gt.to(device, non_blocking=True), text.to( device, non_blocking=True) predicted = model(text) # loss = criterion_bce(predicted, gt) + criterion_dice(predicted, gt) loss = criterion( predicted, gt - text ) # predicts extracted text in white, all others in black epoch_losses.update(loss.item(), len(gt)) optimizer.zero_grad() loss.backward() optimizer.step() _tqdm.set_postfix(loss='{:.6f}'.format(epoch_losses.avg)) _tqdm.update(len(gt)) save_path = './autoencoder_models_2' if not os.path.exists(save_path): os.mkdir(save_path) gt_text = gt - text predicted_mask = text + predicted torch.save( model.state_dict(), os.path.join(save_path, 'autoencoder_{}.pth'.format(epoch + 1 + epoch_plus))) writer.add_scalar('Loss', epoch_losses.avg, epoch + 1 + epoch_plus) writer.add_image('text/text_image_{}'.format(epoch + 1 + epoch_plus), text[0].squeeze(), epoch + 1 + epoch_plus) writer.add_image('gt/gt_image_{}'.format(epoch + 1 + epoch_plus), gt[0].squeeze(), epoch + 1 + epoch_plus) writer.add_image('gt_text/gt_image_{}'.format(epoch + 1 + epoch_plus), gt_text[0].squeeze(), epoch + 1 + epoch_plus) writer.add_image( 'predicted/predicted_image_{}'.format(epoch + 1 + epoch_plus), predicted_mask[0].squeeze(), epoch + 1 + epoch_plus) writer.add_image( 'predicted_text/predicted_image_{}'.format(epoch + 1 + epoch_plus), predicted[0].squeeze(), epoch + 1 + epoch_plus) writer.close()
if __name__ == "__main__": model = SegNet().to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01, betas=(0.9, 0.999)) train_dataset, test_dataset = load_dataset('dataset/TrainingData') scores_avg = evaluate(test_dataset) print(scores_avg) for epoch in range(30): epoch_loss = 0 for i, (x, y) in enumerate(train_dataset.batch(16)): x, y = x.to(device), y.to(device) optimizer.zero_grad() out = model(x) loss = loss_f(out, y) epoch_loss += loss.cpu().item() loss.backward() optimizer.step() print(epoch, epoch_loss / len(train_dataset)) scores_avg = evaluate(test_dataset) print(epoch, scores_avg) print() model_path = f'models/epoch-{epoch}.pth' torch.save(model.state_dict(), model_path)