train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Opion():\n",
    "    \n",
    "    def __init__(self):\n",
    "            \n",
    "        self.dataroot= r'I:\\irregular holes\\paris_eval_gt' #image dataroot\n",
    "        self.maskroot= r'I:\\irregular holes\\testing_mask_dataset'#mask dataroot\n",
    "        self.batchSize= 1   # Need to be set to 1\n",
    "        self.fineSize=256 # image size\n",
    "        self.input_nc=3  # input channel size for first stage\n",
    "        self.input_nc_g=6 # input channel size for second stage\n",
    "        self.output_nc=3# output channel size\n",
    "        self.ngf=64 # inner channel\n",
    "        self.ndf=64# inner channel\n",
    "        self.which_model_netD='basic' # patch discriminator\n",
    "        \n",
    "        self.which_model_netF='feature'# feature patch discriminator\n",
    "        self.which_model_netG='unet_csa'# seconde stage network\n",
    "        self.which_model_netP='unet_256'# first stage network\n",
    "        self.triple_weight=1\n",
    "        self.name='CSA_inpainting'\n",
    "        self.n_layers_D='3' # network depth\n",
    "        self.gpu_ids=[0]\n",
    "        self.model='csa_net'\n",
    "        self.checkpoints_dir=r'.\\checkpoints' #\n",
    "        self.norm='instance'\n",
    "        self.fixed_mask=1\n",
    "        self.use_dropout=False\n",
    "        self.init_type='normal'\n",
    "        self.mask_type='center'\n",
    "        self.lambda_A=100\n",
    "        self.threshold=5/16.0\n",
    "        self.stride=1\n",
    "        self.shift_sz=1 # size of feature patch\n",
    "        self.mask_thred=1\n",
    "        self.bottleneck=512\n",
    "        self.gp_lambda=10.0\n",
    "        self.ncritic=5\n",
    "        self.constrain='MSE'\n",
    "        self.strength=1\n",
    "        self.init_gain=0.02\n",
    "        self.cosis=1\n",
    "        self.gan_type='lsgan'\n",
    "        self.gan_weight=0.2\n",
    "        self.overlap=4\n",
    "        self.skip=0\n",
    "        self.display_freq=1000\n",
    "        self.print_freq=50\n",
    "        self.save_latest_freq=5000\n",
    "        self.save_epoch_freq=2\n",
    "        self.continue_train=False\n",
    "        self.epoch_count=1\n",
    "        self.phase='train'\n",
    "        self.which_epoch=''\n",
    "        self.niter=20\n",
    "        self.niter_decay=100\n",
    "        self.beta1=0.5\n",
    "        self.lr=0.0002\n",
    "        self.lr_policy='lambda'\n",
    "        self.lr_decay_iters=50\n",
    "        self.isTrain=True\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from util.data_load import Data_load\n",
    "from models.models import create_model\n",
    "import torch\n",
    "import os\n",
    "import torchvision\n",
    "from torch.utils import data\n",
    "import torchvision.transforms as transforms\n",
    "opt = Opion()\n",
    "transform_mask = transforms.Compose(\n",
    "    [transforms.Resize((opt.fineSize,opt.fineSize)),\n",
    "     transforms.ToTensor(),\n",
    "    ])\n",
    "transform = transforms.Compose(\n",
    "    [transforms.RandomHorizontalFlip(),\n",
    "     transforms.Resize((opt.fineSize,opt.fineSize)),\n",
    "     transforms.ToTensor(),\n",
    "     transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)])\n",
    "\n",
    "dataset_train = Data_load(opt.dataroot, opt.maskroot, transform, transform_mask)\n",
    "iterator_train = (data.DataLoader(dataset_train, batch_size=opt.batchSize,shuffle=True))\n",
    "print(len(dataset_train))\n",
    "model = create_model(opt)\n",
    "total_steps = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "iter_start_time = time.time()\n",
    "for epoch in range(opt.epoch_count, opt.niter + opt.niter_decay + 1):\n",
    "    \n",
    "    epoch_start_time = time.time()\n",
    "    epoch_iter = 0\n",
    "\n",
    "#     image, mask, gt = [x.cuda() for x in next(iterator_train)]\n",
    "    for image, mask in (iterator_train):\n",
    "        image=image.cuda()\n",
    "        mask=mask.cuda()\n",
    "        mask=mask[0][0]\n",
    "        mask=torch.unsqueeze(mask,0)\n",
    "        mask=torch.unsqueeze(mask,1)\n",
    "        mask=mask.byte()\n",
    "\n",
    "        total_steps += opt.batchSize\n",
    "        epoch_iter += opt.batchSize\n",
    "        model.set_input(image,mask) # it not only sets the input data with mask, but also sets the latent mask.\n",
    "        model.set_gt_latent()\n",
    "        model.optimize_parameters()\n",
    "\n",
    "        if total_steps %opt.display_freq== 0:\n",
    "            real_A,real_B,fake_B=model.get_current_visuals()\n",
    "            #real_A=input, real_B=ground truth fake_b=output\n",
    "            pic = (torch.cat([real_A, real_B,fake_B], dim=0) + 1) / 2.0\n",
    "            torchvision.utils.save_image(pic_1, '%s/Epoch_(%d)_(%dof%d).jpg' % (\n",
    "            save_dir, epoch, total_steps + 1, len(dataset_train)), nrow=2)\n",
    "        if total_steps %1== 0:\n",
    "            errors = model.get_current_errors()\n",
    "            t = (time.time() - iter_start_time) / opt.batchSize\n",
    "            print(errors)\n",
    "\n",
    "    if epoch % opt.save_epoch_freq == 0:\n",
    "        print('saving the model at the end of epoch %d, iters %d' %\n",
    "                (epoch, total_steps))\n",
    "        model.save(epoch)\n",
    "\n",
    "    print('End of epoch %d / %d \\t Time Taken: %d sec' %\n",
    "            (epoch, opt.niter + opt.niter_decay, time.time() - epoch_start_time))\n",
    "\n",
    "    model.update_learning_rate()\n"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "conda-env-pytorch-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}