supervised_train.py

import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import time
from datetime import datetime
from dataset import MyDataset
from optimizer import RangerV2
from utils import dice_coef_2d
from loss import DiceWithBceLoss
import argparse
import os
import torch.optim as optim

parser = argparse.ArgumentParser()

parser.add_argument('--epoch', type=int, default=1201,
                    help='epoch number')
parser.add_argument('--start_iter', type=int, default=1,
                    help='Starting Epoch')
parser.add_argument('--lr', type=float, default=5e-4,
                    help='learning rate')
parser.add_argument('--decay_rate', type=float, default=0.1,
                    help='decay rate of learning rate')
parser.add_argument('--decay_start_epoch', type=int, default=800,
                    help='when to start to decay learning rate')
parser.add_argument('--decay_epoch', type=int, default=30,
                    help='every n epochs decay learning rate')
parser.add_argument('--batchsize', type=int, default=4,
                    help='training batch size')
parser.add_argument('--data_augmentation', type=bool, default=False)

parser.add_argument('--num_workers', type=int, default=0,
                    help='number of workers in dataloader. In windows, set num_workers=0')
parser.add_argument('--train_img_path', type=str, default=r'.\Train-Image',
                    help='images path for training')
parser.add_argument('--train_msk_path', type=str, default=r'.\Train-Mask',
                    help='images mask path for training')
parser.add_argument('--optimizer_type', type=str, default='Ranger',
                    help='type of optimizer')
parser.add_argument('--pretrained', type=bool, default=False)

parser.add_argument('--pretrained_model', type=str, default='',
                    help='pretrained base model')
parser.add_argument('--save_start_epoch', type=int, default=200,
                    help='starting to save model epoch ')
parser.add_argument('--snapshots', type=int, default=30,
                    help='every n epochs save a model')

parser.add_argument('--save_folder', type=str, default='./',
                    help='Location to save checkpoint models')


opt = parser.parse_args()
if opt.data_augmentation:
    train_transform = transforms.Compose([transforms.RandomHorizontalFlip(),
                                          transforms.RandomVerticalFlip(),
                                          transforms.RandomRotation(90),
                                          transforms.ToTensor()
                                          ])
else:
    train_transform = transforms.Compose([transforms.ToTensor()])

train_dataset = MyDataset(opt.train_img_path, opt.train_msk_path, train_transform)
train_dataloader = DataLoader(train_dataset, batch_size=opt.batchsize, shuffle=True, num_workers=opt.num_workers)

epo_num = opt.epoch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

from model import CARes_Unet
model = CARes_Unet()



if opt.pretrained:
    if os.path.exists(opt.pretrained_model):
        model.load_state_dict(torch.load(opt.pretrained_model, map_location=lambda storage, loc: storage))
        print('Pre-trained model is loaded.')
    else:
        print("Model Not Found")
        exit(-1)
model = model.to(device)

criterion = DiceWithBceLoss().to(device)
if opt.optimizer_type == 'Ranger':
    optimizer = RangerV2(model.parameters(), lr=opt.lr)
elif opt.optimizer_type == 'Adam':
    optimizer = optim.Adam(model.parameters(), lr=opt.lr)
elif opt.optimizer_type == 'SGD':
    optimizer = optim.SGD(model.parameters(), lr=opt.lr)
print('optimizer:', opt.optimizer_type)

prev_time = datetime.now()

for epo in range(opt.start_iter, epo_num):
    train_loss = 0
    train_dice = 0
    model.train()
    for index, (img, img_msk) in enumerate(train_dataloader):
        img = img.to(device)
        img_msk = img_msk.to(device)

        optimizer.zero_grad()
        output = model(img)
        loss = criterion(output, img_msk)
        dice = dice_coef_2d(output, img_msk)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        train_dice += dice.item()

    # compute time
    cur_time = datetime.now()
    h, remainder = divmod((cur_time - prev_time).seconds, 3600)
    m, s = divmod(remainder, 60)
    time_str = "Time %02d:%02d:%02d" % (h, m, s)
    prev_time = cur_time

    print('epoch: %f, train loss = %f, train dice = %f ,%s'
          % (epo, train_loss / len(train_dataloader), train_dice / len(train_dataloader), time_str))
    if epo % opt.decay_epoch == 0 and epo > opt.decay_start_epoch:
        for param_group in optimizer.param_groups:
            param_group['lr'] *= opt.decay_rate
        print('Learning rate decay: lr={}'.format(optimizer.param_groups[0]['lr']))

    if epo >opt.save_start_epoch and epo %opt.snapshots==0:
        save_path = opt.save_folder + 'supervised_' + opt.model_type + '_epoch_{}.pth'.format(epo)
        torch.save(model.state_dict(), save_path)
        print("model_copy is saved !")