import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data
from torch.optim.lr_scheduler import LambdaLR
import torch.nn.functional as F
from torchvision import transforms as T
from data_loader import GetLoader, ForeverDataIterator
from torch.autograd import Variable
from model import ILA
from gan import AugNet
import os
import datetime
import argparse
from logger import Logger
from score import ScoreDA
print(1)
def load_train_data():
# transform
# img_transform = T.Compose([
# T.Resize(256),
# T.CenterCrop(224),
# # T.RandomResizedCrop(224),
# T.RandomHorizontalFlip(),
# T.ToTensor(),
# T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# ])
img_transform = T.Compose([
T.Resize(128),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=0.5, std=0.5)
])
# source
dataset_src = GetLoader(
data_root=args.img_src,
data_list=args.lab_src,
transform=img_transform
)
dataloader_src = torch.utils.data.DataLoader(
dataset=dataset_src,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=0
)
# target
dataset_tar = GetLoader(
data_root=args.img_tar,
data_list=args.lab_tar,
transform=img_transform
)
dataloader_tar = torch.utils.data.DataLoader(
dataset=dataset_tar,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=0
)
return dataloader_src, dataloader_tar
def load_test_data():
# transform
img_transform = T.Compose([
T.Resize(128),
T.ToTensor(),
T.Normalize(mean=0.5, std=0.5)
])
dataset_tar = GetLoader(
data_root=args.img_tar,
data_list=args.lab_tar,
transform=img_transform
)
dataloader_tar = torch.utils.data.DataLoader(
dataset=dataset_tar,
batch_size=args.batch_size,
shuffle=True,
drop_last=True,
num_workers=0
)
return dataloader_tar
def test(net):
data_loader = load_test_data()
score_test = ScoreDA()
dataloader_iter_test = iter(data_loader)
net.eval()
with torch.no_grad():
for img, lab in dataloader_iter_test:
img, lab = img.to(args.device), lab.to(args.device)
rlt, _ = net(img)
score_test.updateDA(lab, rlt, 'tar')
score_test.showTar()
print(score_test.show_acc(mode='src'), score_test.show_acc(mode='tar'))
def find_pseudo(rlt, num):
rlt = rlt.detach().cpu().numpy()
a = np.split(rlt, num+1)
for i in range(len(a)):
a[i] = np.argmax(a[i], axis=1)
a = np.array(a)
pseudo = list()
for i in range(0, args.batch_size):
b = np.argmax(np.bincount(a[:, i]))
pseudo.append(b)
pseudo = torch.tensor(pseudo)
# print(pseudo)
pseudo_label = list()
for i in range(num+1):
pseudo_label.append(pseudo)
pseudo_label = torch.cat(pseudo_label, dim=0).cuda()
# print(pseudo_label)
return pseudo_label
def main():
print(args)
# logger
logger = Logger(args)
# score
score = ScoreDA()
# model
net = ILA(num_cls=args.num_cls, iterations=args.iteration).to(args.device)
gan = AugNet(1).to(args.device)
# optimizer
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-3, nesterov=True)
optimizer_gan = optim.SGD(list(gan.parameters()) + list(net.parameters()), lr=args.lr, momentum=0.9,
weight_decay=1e-3, nesterov=True)
lr_scheduler = LambdaLR(optimizer_gan, lambda x: args.lr * (1. + args.lr_gamma * float(x)) ** (-args.lr_decay))
# data iter
dataloader_src, dataloader_tar = load_train_data()
print('length:', 'src:', len(dataloader_src), 'tar:', len(dataloader_tar))
dataiter_src, dataiter_tar = ForeverDataIterator(dataloader_src), ForeverDataIterator(dataloader_tar)
# loss
# training
for step in range(args.iteration):
net.train()
gan.train()
# data
img_src, lab_src = next(dataiter_src)
img_tar, lab_tar = next(dataiter_tar)
img_src, lab_src = img_src.to(args.device), lab_src.to(args.device)
img_tar, lab_tar = img_tar.to(args.device), lab_tar.to(args.device)
# domainlab_src = torch.zeros(args.batch_size).long().to(args.device)
# domainlab_tar = torch.ones(args.batch_size).long().to(args.device)
# model
inputs_maxs = gan(img_src, "gan")
data_gen = list()
num_gen = 0
for inputs_max in inputs_maxs:
inputs_max = torch.sigmoid(inputs_max)
#inputs_max = tran(torch.sigmoid(inputs_max))
inputs_max = inputs_max * 0.4 + img_src * 0.6
data_gen.append(inputs_max)
num_gen += 1
data_gen = torch.cat(data_gen, dim=0)
data_aug = torch.cat([img_src, data_gen], dim=0)
domian_lable = list()
lab_aug_src = list()
for i in range(num_gen+1):
domian_lable.append(Variable((torch.ones(args.batch_size)).long().cuda()))
lab_aug_src.append(Variable(lab_src.detach().cuda()))
lab_aug_src = torch.cat(lab_aug_src, dim=0)
# source and augment
domian_lable = torch.cat(domian_lable, dim=0)
cls_pre, domain_pre = net(data_aug)
loss_domain = F.cross_entropy(cls_pre, lab_aug_src) + F.cross_entropy(domain_pre, domian_lable)
score.updateDA(lab_src, cls_pre[0:32], 'src')
optimizer.zero_grad()
loss_domain.backward()
optimizer.step()
# target and augment
inputs_maxs_tar = gan(img_tar, "gan")
data_gen_tar = list()
num_gen = 0
for inputs_max in inputs_maxs_tar:
inputs_max = torch.sigmoid(inputs_max)
# inputs_max = tran(torch.sigmoid(inputs_max))
inputs_max = inputs_max * 0.4 + img_tar * 0.6
data_gen_tar.append(inputs_max)
num_gen += 1
data_gen_tar = torch.cat(data_gen_tar, dim=0)
data_aug_tar = torch.cat([img_tar, data_gen_tar], dim=0)
domian_lab_tar = list()
for i in range(num_gen+1):
domian_lab_tar.append(Variable((torch.zeros(args.batch_size)).long().cuda()))
domian_lab_tar = torch.cat(domian_lab_tar, dim=0)
rlt_tar, domain_tar = net(data_aug_tar)
lab_pseudo_tar = find_pseudo(rlt_tar, num_gen)
loss = F.cross_entropy(domain_tar, domian_lab_tar) + 0.2 * F.cross_entropy(rlt_tar, lab_pseudo_tar)
score.updateDA(lab_tar, rlt_tar[0:32], 'cls')
optimizer.zero_grad()
loss.backward()
optimizer.step()
# learning rate
lr_scheduler.step()
# print and save
if step % args.print_freq == 0:
print(step, 'src:', score.show_acc(mode='src'), 'tar:', score.show_acc(mode='tar'))
if step % args.test_freq == 0:
print('src:', score.show_acc(mode='src'), 'tar:', score.show_acc(mode='tar'))
score.reset()
# test(net, args)
torch.save(net, '{}/{}-{}/{}-{}-step-{}.pth'.format(args.log, args.name, args.time, args.name, args.time, step))
logger.close()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--img_src', type=str, default='./datasets/HanDanSteel', dest='img_src')
parser.add_argument('--lab_src', type=str, default='./datasets/HanDanSteel/HanDanSteel-low.txt', dest='lab_src')
parser.add_argument('--img_tar', type=str, default='./datasets/HanDanSteel', dest='img_tar')
parser.add_argument('--lab_tar', type=str, default='./datasets/HanDanSteel/HanDanSteel-middle.txt', dest='lab_tar')
parser.add_argument('--name', type=str, default='HanDanSteel-low2middle', dest='name')
parser.add_argument('--iteration', type=int, default=10000, dest='iteration', help='Number of iteration')
parser.add_argument('--num_cls', type=int, default=6, dest='num_cls')
parser.add_argument('--lr', type=float, default=1e-2, dest='lr', help='learning rate')
parser.add_argument('--lr_gamma', type=float, default=0.0003, dest='lr_gamma')
parser.add_argument('--lr_decay', type=float, default=0.75, dest='lr_decay')
parser.add_argument('--batch_size', type=int, default=32, dest='batch_size')
parser.add_argument('--time', type=str, default=datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"), dest='time')
parser.add_argument('--device', type=str, dest='device',
default=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))
parser.add_argument('--seed', type=int, default=100, dest='seed')
parser.add_argument('--log', type=str, default='./log', dest='log')
parser.add_argument('--print_freq', type=int, default=1, dest='print_freq')
parser.add_argument('--test_freq', type=int, default=500, dest='test_freq')
args = parser.parse_args()
torch.random.manual_seed(args.seed)
main()
