wogong
7 years ago
17 changed files with 793 additions and 0 deletions
@ -0,0 +1,27 @@ |
|||
# PyTorch-DANN |
|||
|
|||
A pytorch implementation for paper *[Unsupervised Domain Adaptation by Backpropagation](http://sites.skoltech.ru/compvision/projects/grl/)* |
|||
|
|||
InProceedings (icml2015-ganin15) |
|||
Ganin, Y. & Lempitsky, V. |
|||
Unsupervised Domain Adaptation by Backpropagation |
|||
Proceedings of the 32nd International Conference on Machine Learning, 2015 |
|||
|
|||
## Environment |
|||
|
|||
- Python 2.7 |
|||
- PyTorch 0.3.1 |
|||
|
|||
## Result |
|||
|
|||
results of the default `params.py` |
|||
|
|||
| | MNIST (Source) | USPS (Target) | |
|||
| :--------------------------------: | :------------: | :-----------: | |
|||
| Source Classifier | 99.140000% | 83.978495% | |
|||
| DANN | | 97.634409% | |
|||
|
|||
## Credit |
|||
|
|||
- <https://github.com/fungtion/DANN> |
|||
- <https://github.com/corenel/torchsharp> |
|||
@ -0,0 +1,96 @@ |
|||
"""Train dann.""" |
|||
|
|||
import torch |
|||
import torch.nn as nn |
|||
import torch.optim as optim |
|||
|
|||
import params |
|||
from utils import make_variable, save_model |
|||
import numpy as np |
|||
from core.test import eval |
|||
|
|||
import torch.backends.cudnn as cudnn |
|||
cudnn.benchmark = True |
|||
|
|||
def train_dann(dann, src_data_loader, tgt_data_loader, tgt_data_loader_eval): |
|||
"""Train dann.""" |
|||
#################### |
|||
# 1. setup network # |
|||
#################### |
|||
|
|||
# set train state for Dropout and BN layers |
|||
dann.train() |
|||
|
|||
# setup criterion and optimizer |
|||
optimizer = optim.Adam(dann.parameters(), lr=params.lr) |
|||
|
|||
criterion = nn.NLLLoss() |
|||
|
|||
for p in dann.parameters(): |
|||
p.requires_grad = True |
|||
|
|||
#################### |
|||
# 2. train network # |
|||
#################### |
|||
|
|||
# prepare domain label |
|||
label_src = make_variable(torch.zeros(params.batch_size).long()) # source 0 |
|||
label_tgt = make_variable(torch.ones(params.batch_size).long()) # target 1 |
|||
|
|||
for epoch in range(params.num_epochs): |
|||
# zip source and target data pair |
|||
len_dataloader = min(len(src_data_loader), len(tgt_data_loader)) |
|||
data_zip = enumerate(zip(src_data_loader, tgt_data_loader)) |
|||
for step, ((images_src, class_src), (images_tgt, _)) in data_zip: |
|||
|
|||
p = float(step + epoch * len_dataloader) / params.num_epochs / len_dataloader |
|||
alpha = 2. / (1. + np.exp(-10 * p)) - 1 |
|||
|
|||
# make images variable |
|||
class_src = make_variable(class_src) |
|||
images_src = make_variable(images_src) |
|||
images_tgt = make_variable(images_tgt) |
|||
|
|||
# zero gradients for optimizer |
|||
optimizer.zero_grad() |
|||
|
|||
# train on source domain |
|||
src_class_output, src_domain_output = dann(input_data=images_src, alpha=alpha) |
|||
src_loss_class = criterion(src_class_output, class_src) |
|||
src_loss_domain = criterion(src_domain_output, label_src) |
|||
|
|||
# train on target domain |
|||
_, tgt_domain_output = dann(input_data=images_tgt, alpha=alpha) |
|||
tgt_loss_domain = criterion(tgt_domain_output, label_tgt) |
|||
|
|||
loss = src_loss_class + src_loss_domain + tgt_loss_domain |
|||
|
|||
# optimize dann |
|||
loss.backward() |
|||
optimizer.step() |
|||
|
|||
# print step info |
|||
if ((step + 1) % params.log_step == 0): |
|||
print("Epoch [{}/{}] Step [{}/{}]: src_loss_class={}, src_loss_domain={}, tgt_loss_domain={}, loss={}" |
|||
.format(epoch + 1, |
|||
params.num_epochs, |
|||
step + 1, |
|||
len_dataloader, |
|||
src_loss_class.data[0], |
|||
src_loss_domain.data[0], |
|||
tgt_loss_domain.data[0], |
|||
loss.data[0])) |
|||
|
|||
# eval model on test set |
|||
if ((epoch + 1) % params.eval_step == 0): |
|||
eval(dann, tgt_data_loader_eval) |
|||
dann.train() |
|||
|
|||
# save model parameters |
|||
if ((epoch + 1) % params.save_step == 0): |
|||
save_model(dann, params.src_dataset + '-' + params.tgt_dataset + "-dann-{}.pt".format(epoch + 1)) |
|||
|
|||
# save final model |
|||
save_model(dann, params.src_dataset + '-' + params.tgt_dataset + "-dann-final.pt") |
|||
|
|||
return dann |
|||
@ -0,0 +1,65 @@ |
|||
"""Train classifier for source dataset.""" |
|||
|
|||
import torch.nn as nn |
|||
import torch.optim as optim |
|||
|
|||
import params |
|||
from utils import make_variable, save_model |
|||
from core.test import eval_src |
|||
|
|||
def train_src(model, data_loader): |
|||
"""Train classifier for source domain.""" |
|||
#################### |
|||
# 1. setup network # |
|||
#################### |
|||
|
|||
# set train state for Dropout and BN layers |
|||
model.train() |
|||
|
|||
# setup criterion and optimizer |
|||
optimizer = optim.Adam(model.parameters(), lr=params.lr) |
|||
loss_class = nn.NLLLoss() |
|||
|
|||
#################### |
|||
# 2. train network # |
|||
#################### |
|||
|
|||
for epoch in range(params.num_epochs_src): |
|||
for step, (images, labels) in enumerate(data_loader): |
|||
# make images and labels variable |
|||
images = make_variable(images) |
|||
labels = make_variable(labels.squeeze_()) |
|||
|
|||
# zero gradients for optimizer |
|||
optimizer.zero_grad() |
|||
|
|||
# compute loss for critic |
|||
preds = model(images) |
|||
loss = loss_class(preds, labels) |
|||
|
|||
# optimize source classifier |
|||
loss.backward() |
|||
optimizer.step() |
|||
|
|||
# print step info |
|||
if ((step + 1) % params.log_step_src == 0): |
|||
print("Epoch [{}/{}] Step [{}/{}]: loss={}" |
|||
.format(epoch + 1, |
|||
params.num_epochs_src, |
|||
step + 1, |
|||
len(data_loader), |
|||
loss.data[0])) |
|||
|
|||
# eval model on test set |
|||
if ((epoch + 1) % params.eval_step_src == 0): |
|||
eval_src(model, data_loader) |
|||
model.train() |
|||
|
|||
# save model parameters |
|||
if ((epoch + 1) % params.save_step_src == 0): |
|||
save_model(model, params.src_dataset + "-source-classifier-{}.pt".format(epoch + 1)) |
|||
|
|||
# save final model |
|||
save_model(model, params.src_dataset + "-source-classifier-final.pt") |
|||
|
|||
return model |
|||
@ -0,0 +1,159 @@ |
|||
import os |
|||
import torch.backends.cudnn as cudnn |
|||
import torch.utils.data |
|||
from torch.autograd import Variable |
|||
import torch.nn as nn |
|||
|
|||
import params |
|||
from utils import make_variable |
|||
|
|||
def test(dataloader, epoch): |
|||
|
|||
if torch.cuda.is_available(): |
|||
cuda = True |
|||
cudnn.benchmark = True |
|||
|
|||
""" training """ |
|||
|
|||
my_net = torch.load(os.path.join( |
|||
params.model_root, params.src_dataset + '_' + params.tgt_dataset + '_model_epoch_' + str(epoch) + '.pth' |
|||
)) |
|||
my_net = my_net.eval() |
|||
|
|||
if cuda: |
|||
my_net = my_net.cuda() |
|||
|
|||
len_dataloader = len(dataloader) |
|||
data_target_iter = iter(dataloader) |
|||
|
|||
i = 0 |
|||
n_total = 0 |
|||
n_correct = 0 |
|||
|
|||
while i < len_dataloader: |
|||
|
|||
# test model using target data |
|||
data_target = data_target_iter.next() |
|||
t_img, t_label = data_target |
|||
|
|||
batch_size = len(t_label) |
|||
|
|||
input_img = torch.FloatTensor(batch_size, 3, params.image_size, params.image_size) |
|||
class_label = torch.LongTensor(batch_size) |
|||
|
|||
if cuda: |
|||
t_img = t_img.cuda() |
|||
t_label = t_label.cuda() |
|||
input_img = input_img.cuda() |
|||
class_label = class_label.cuda() |
|||
|
|||
input_img.resize_as_(t_img).copy_(t_img) |
|||
class_label.resize_as_(t_label).copy_(t_label) |
|||
inputv_img = Variable(input_img) |
|||
classv_label = Variable(class_label) |
|||
|
|||
class_output, _ = my_net(input_data=inputv_img, alpha=params.alpha) |
|||
pred = class_output.data.max(1, keepdim=True)[1] |
|||
n_correct += pred.eq(classv_label.data.view_as(pred)).cpu().sum() |
|||
n_total += batch_size |
|||
|
|||
i += 1 |
|||
|
|||
accu = n_correct * 1.0 / n_total |
|||
|
|||
print 'epoch: %d, accuracy: %f' % (epoch, accu) |
|||
|
|||
|
|||
def test_from_save(model, saved_model, data_loader): |
|||
"""Evaluate classifier for source domain.""" |
|||
# set eval state for Dropout and BN layers |
|||
classifier = model.load_state_dict(torch.load(saved_model)) |
|||
classifier.eval() |
|||
|
|||
# init loss and accuracy |
|||
loss = 0.0 |
|||
acc = 0.0 |
|||
|
|||
# set loss function |
|||
criterion = nn.NLLLoss() |
|||
|
|||
# evaluate network |
|||
for (images, labels) in data_loader: |
|||
images = make_variable(images, volatile=True) |
|||
labels = make_variable(labels) #labels = labels.squeeze(1) |
|||
preds = classifier(images) |
|||
|
|||
criterion(preds, labels) |
|||
|
|||
loss += criterion(preds, labels).data[0] |
|||
|
|||
pred_cls = preds.data.max(1)[1] |
|||
acc += pred_cls.eq(labels.data).cpu().sum() |
|||
|
|||
loss /= len(data_loader) |
|||
acc /= len(data_loader.dataset) |
|||
|
|||
print("Avg Loss = {}, Avg Accuracy = {:2%}".format(loss, acc)) |
|||
|
|||
|
|||
def eval(model, data_loader): |
|||
"""Evaluate model for dataset.""" |
|||
# set eval state for Dropout and BN layers |
|||
model.eval() |
|||
|
|||
# init loss and accuracy |
|||
loss = 0.0 |
|||
acc = 0.0 |
|||
|
|||
# set loss function |
|||
criterion = nn.NLLLoss() |
|||
|
|||
# evaluate network |
|||
for (images, labels) in data_loader: |
|||
images = make_variable(images, volatile=True) |
|||
labels = make_variable(labels) #labels = labels.squeeze(1) |
|||
preds, _ = model(images, alpha=0) |
|||
|
|||
criterion(preds, labels) |
|||
|
|||
loss += criterion(preds, labels).data[0] |
|||
|
|||
pred_cls = preds.data.max(1)[1] |
|||
acc += pred_cls.eq(labels.data).cpu().sum() |
|||
|
|||
loss /= len(data_loader) |
|||
acc /= len(data_loader.dataset) |
|||
|
|||
print("Avg Loss = {}, Avg Accuracy = {:2%}".format(loss, acc)) |
|||
|
|||
|
|||
def eval_src(model, data_loader): |
|||
"""Evaluate classifier for source domain.""" |
|||
# set eval state for Dropout and BN layers |
|||
model.eval() |
|||
|
|||
# init loss and accuracy |
|||
loss = 0.0 |
|||
acc = 0.0 |
|||
|
|||
# set loss function |
|||
criterion = nn.NLLLoss() |
|||
|
|||
# evaluate network |
|||
for (images, labels) in data_loader: |
|||
images = make_variable(images, volatile=True) |
|||
labels = make_variable(labels) #labels = labels.squeeze(1) |
|||
preds = model(images) |
|||
|
|||
criterion(preds, labels) |
|||
|
|||
loss += criterion(preds, labels).data[0] |
|||
|
|||
pred_cls = preds.data.max(1)[1] |
|||
acc += pred_cls.eq(labels.data).cpu().sum() |
|||
|
|||
|
|||
loss /= len(data_loader) |
|||
acc /= len(data_loader.dataset) |
|||
|
|||
print("Avg Loss = {}, Avg Accuracy = {:2%}".format(loss, acc)) |
|||
@ -0,0 +1,5 @@ |
|||
from .mnist import get_mnist |
|||
from .mnistm import get_mnistm |
|||
from .svhn import get_svhn |
|||
|
|||
__all__ = (get_mnist, get_svhn, get_mnistm) |
|||
@ -0,0 +1,31 @@ |
|||
"""Dataset setting and data loader for MNIST.""" |
|||
|
|||
|
|||
import torch |
|||
from torchvision import datasets, transforms |
|||
import os |
|||
|
|||
import params |
|||
|
|||
def get_mnist(train): |
|||
"""Get MNIST datasets loader.""" |
|||
# image pre-processing |
|||
pre_process = transforms.Compose([transforms.ToTensor(), |
|||
transforms.Normalize( |
|||
mean=params.dataset_mean, |
|||
std=params.dataset_std)]) |
|||
|
|||
# datasets and data loader |
|||
mnist_dataset = datasets.MNIST(root=os.path.join(params.dataset_root,'mnist'), |
|||
train=train, |
|||
transform=pre_process, |
|||
download=False) |
|||
|
|||
|
|||
mnist_data_loader = torch.utils.data.DataLoader( |
|||
dataset=mnist_dataset, |
|||
batch_size=params.batch_size, |
|||
shuffle=True, |
|||
drop_last=True) |
|||
|
|||
return mnist_data_loader |
|||
@ -0,0 +1,70 @@ |
|||
"""Dataset setting and data loader for MNIST_M.""" |
|||
|
|||
import torch |
|||
from torchvision import datasets, transforms |
|||
import torch.utils.data as data |
|||
from PIL import Image |
|||
import os |
|||
import params |
|||
|
|||
class GetLoader(data.Dataset): |
|||
def __init__(self, data_root, data_list, transform=None): |
|||
self.root = data_root |
|||
self.transform = transform |
|||
|
|||
f = open(data_list, 'r') |
|||
data_list = f.readlines() |
|||
f.close() |
|||
|
|||
self.n_data = len(data_list) |
|||
|
|||
self.img_paths = [] |
|||
self.img_labels = [] |
|||
|
|||
for data in data_list: |
|||
self.img_paths.append(data[:-3]) |
|||
self.img_labels.append(data[-2]) |
|||
|
|||
def __getitem__(self, item): |
|||
img_paths, labels = self.img_paths[item], self.img_labels[item] |
|||
imgs = Image.open(os.path.join(self.root, img_paths)).convert('RGB') |
|||
|
|||
if self.transform is not None: |
|||
imgs = self.transform(imgs) |
|||
labels = int(labels) |
|||
|
|||
return imgs, labels |
|||
|
|||
def __len__(self): |
|||
return self.n_data |
|||
|
|||
def get_mnistm(train): |
|||
"""Get MNISTM datasets loader.""" |
|||
# image pre-processing |
|||
pre_process = transforms.Compose([transforms.Resize(params.image_size), |
|||
transforms.ToTensor(), |
|||
transforms.Normalize( |
|||
mean=params.dataset_mean, |
|||
std=params.dataset_std)]) |
|||
|
|||
# datasets and data_loader |
|||
if train: |
|||
train_list = os.path.join(params.dataset_root, 'mnist_m','mnist_m_train_labels.txt') |
|||
mnistm_dataset = GetLoader( |
|||
data_root=os.path.join(params.dataset_root, 'mnist_m', 'mnist_m_train'), |
|||
data_list=train_list, |
|||
transform=pre_process) |
|||
else: |
|||
train_list = os.path.join(params.dataset_root, 'mnist_m', 'mnist_m_test_labels.txt') |
|||
mnistm_dataset = GetLoader( |
|||
data_root=os.path.join(params.dataset_root, 'mnist_m', 'mnist_m_test'), |
|||
data_list=train_list, |
|||
transform=pre_process) |
|||
|
|||
mnistm_dataloader = torch.utils.data.DataLoader( |
|||
dataset=mnistm_dataset, |
|||
batch_size=params.batch_size, |
|||
shuffle=True, |
|||
num_workers=8) |
|||
|
|||
return mnistm_dataloader |
|||
@ -0,0 +1,39 @@ |
|||
"""Dataset setting and data loader for SVHN.""" |
|||
|
|||
|
|||
import torch |
|||
from torchvision import datasets, transforms |
|||
import os |
|||
|
|||
import params |
|||
|
|||
|
|||
def get_svhn(train): |
|||
"""Get SVHN datasets loader.""" |
|||
# image pre-processing |
|||
pre_process = transforms.Compose([transforms.Grayscale(), |
|||
transforms.Resize(params.image_size), |
|||
transforms.ToTensor(), |
|||
transforms.Normalize( |
|||
mean=params.dataset_mean, |
|||
std=params.dataset_std)]) |
|||
|
|||
# datasets and data loader |
|||
if train: |
|||
svhn_dataset = datasets.SVHN(root=os.path.join(params.dataset_root,'svhn'), |
|||
split='train', |
|||
transform=pre_process, |
|||
download=True) |
|||
else: |
|||
svhn_dataset = datasets.SVHN(root=os.path.join(params.dataset_root,'svhn'), |
|||
split='test', |
|||
transform=pre_process, |
|||
download=True) |
|||
|
|||
svhn_data_loader = torch.utils.data.DataLoader( |
|||
dataset=svhn_dataset, |
|||
batch_size=params.batch_size, |
|||
shuffle=True, |
|||
drop_last=True) |
|||
|
|||
return svhn_data_loader |
|||
@ -0,0 +1,50 @@ |
|||
from models.model import CNNModel |
|||
from models.classifier import Classifier |
|||
|
|||
from core.dann import train_dann |
|||
from core.test import eval, eval_src |
|||
from core.pretrain import train_src |
|||
|
|||
import params |
|||
from utils import get_data_loader, init_model, init_random_seed |
|||
|
|||
# init random seed |
|||
init_random_seed(params.manual_seed) |
|||
|
|||
# load dataset |
|||
src_data_loader = get_data_loader(params.src_dataset) |
|||
src_data_loader_eval = get_data_loader(params.src_dataset, train=False) |
|||
tgt_data_loader = get_data_loader(params.tgt_dataset) |
|||
tgt_data_loader_eval = get_data_loader(params.tgt_dataset, train=False) |
|||
|
|||
# load source classifier |
|||
src_classifier = init_model(net=Classifier(), restore=params.src_classifier_restore) |
|||
|
|||
# train source model |
|||
print("=== Training classifier for source domain ===") |
|||
|
|||
if not (src_classifier.restored and params.src_model_trained): |
|||
src_classifier = train_src(src_classifier, src_data_loader) |
|||
|
|||
# eval source model on both source and target domain |
|||
print("=== Evaluating source classifier for source domain ===") |
|||
eval_src(src_classifier, src_data_loader_eval) |
|||
print("=== Evaluating source classifier for target domain ===") |
|||
eval_src(src_classifier, tgt_data_loader_eval) |
|||
|
|||
# load dann model |
|||
dann = init_model(net=CNNModel(), restore=params.dann_restore) |
|||
|
|||
# train dann model |
|||
print("=== Training dann model ===") |
|||
|
|||
if not (dann.restored and params.dann_restore): |
|||
dann = train_dann(dann, src_data_loader, tgt_data_loader, tgt_data_loader_eval) |
|||
|
|||
# eval dann model |
|||
print("=== Evaluating dann for source domain ===") |
|||
eval(dann, src_data_loader_eval) |
|||
print("=== Evaluating dann for target domain ===") |
|||
eval(dann, tgt_data_loader_eval) |
|||
|
|||
print('done') |
|||
@ -0,0 +1,39 @@ |
|||
"""Classifier for source domain""" |
|||
|
|||
import torch.nn as nn |
|||
|
|||
class Classifier(nn.Module): |
|||
|
|||
def __init__(self): |
|||
super(Classifier, self).__init__() |
|||
self.restored = False |
|||
|
|||
self.feature = nn.Sequential() |
|||
self.feature.add_module('f_conv1', nn.Conv2d(1, 64, kernel_size=5)) |
|||
self.feature.add_module('f_bn1', nn.BatchNorm2d(64)) |
|||
self.feature.add_module('f_pool1', nn.MaxPool2d(2)) |
|||
self.feature.add_module('f_relu1', nn.ReLU(True)) |
|||
self.feature.add_module('f_conv2', nn.Conv2d(64, 50, kernel_size=5)) |
|||
self.feature.add_module('f_bn2', nn.BatchNorm2d(50)) |
|||
self.feature.add_module('f_drop1', nn.Dropout2d()) |
|||
self.feature.add_module('f_pool2', nn.MaxPool2d(2)) |
|||
self.feature.add_module('f_relu2', nn.ReLU(True)) |
|||
|
|||
self.class_classifier = nn.Sequential() |
|||
self.class_classifier.add_module('c_fc1', nn.Linear(50 * 4 * 4, 100)) |
|||
self.class_classifier.add_module('c_bn1', nn.BatchNorm2d(100)) |
|||
self.class_classifier.add_module('c_relu1', nn.ReLU(True)) |
|||
self.class_classifier.add_module('c_drop1', nn.Dropout2d()) |
|||
self.class_classifier.add_module('c_fc2', nn.Linear(100, 100)) |
|||
self.class_classifier.add_module('c_bn2', nn.BatchNorm2d(100)) |
|||
self.class_classifier.add_module('c_relu2', nn.ReLU(True)) |
|||
self.class_classifier.add_module('c_fc3', nn.Linear(100, 10)) |
|||
self.class_classifier.add_module('c_softmax', nn.LogSoftmax(dim=1)) |
|||
|
|||
def forward(self, input_data): |
|||
input_data = input_data.expand(input_data.data.shape[0], 1, 28, 28) |
|||
feature = self.feature(input_data) |
|||
feature = feature.view(-1, 50 * 4 * 4) |
|||
class_output = self.class_classifier(feature) |
|||
|
|||
return class_output |
|||
@ -0,0 +1,18 @@ |
|||
from torch.autograd import Function |
|||
|
|||
|
|||
class ReverseLayerF(Function): |
|||
|
|||
@staticmethod |
|||
def forward(ctx, x, alpha): |
|||
ctx.alpha = alpha |
|||
|
|||
return x.view_as(x) |
|||
|
|||
@staticmethod |
|||
def backward(ctx, grad_output): |
|||
output = grad_output.neg() * ctx.alpha |
|||
|
|||
return output, None |
|||
|
|||
|
|||
@ -0,0 +1,50 @@ |
|||
"""DANN model.""" |
|||
|
|||
import torch.nn as nn |
|||
from functions import ReverseLayerF |
|||
|
|||
|
|||
class CNNModel(nn.Module): |
|||
|
|||
def __init__(self): |
|||
super(CNNModel, self).__init__() |
|||
self.restored = False |
|||
|
|||
self.feature = nn.Sequential() |
|||
self.feature.add_module('f_conv1', nn.Conv2d(1, 64, kernel_size=5)) |
|||
self.feature.add_module('f_bn1', nn.BatchNorm2d(64)) |
|||
self.feature.add_module('f_pool1', nn.MaxPool2d(2)) |
|||
self.feature.add_module('f_relu1', nn.ReLU(True)) |
|||
self.feature.add_module('f_conv2', nn.Conv2d(64, 50, kernel_size=5)) |
|||
self.feature.add_module('f_bn2', nn.BatchNorm2d(50)) |
|||
self.feature.add_module('f_drop1', nn.Dropout2d()) |
|||
self.feature.add_module('f_pool2', nn.MaxPool2d(2)) |
|||
self.feature.add_module('f_relu2', nn.ReLU(True)) |
|||
|
|||
self.class_classifier = nn.Sequential() |
|||
self.class_classifier.add_module('c_fc1', nn.Linear(50 * 4 * 4, 100)) |
|||
self.class_classifier.add_module('c_bn1', nn.BatchNorm2d(100)) |
|||
self.class_classifier.add_module('c_relu1', nn.ReLU(True)) |
|||
self.class_classifier.add_module('c_drop1', nn.Dropout2d()) |
|||
self.class_classifier.add_module('c_fc2', nn.Linear(100, 100)) |
|||
self.class_classifier.add_module('c_bn2', nn.BatchNorm2d(100)) |
|||
self.class_classifier.add_module('c_relu2', nn.ReLU(True)) |
|||
self.class_classifier.add_module('c_fc3', nn.Linear(100, 10)) |
|||
self.class_classifier.add_module('c_softmax', nn.LogSoftmax(dim=1)) |
|||
|
|||
self.domain_classifier = nn.Sequential() |
|||
self.domain_classifier.add_module('d_fc1', nn.Linear(50 * 4 * 4, 100)) |
|||
self.domain_classifier.add_module('d_bn1', nn.BatchNorm2d(100)) |
|||
self.domain_classifier.add_module('d_relu1', nn.ReLU(True)) |
|||
self.domain_classifier.add_module('d_fc2', nn.Linear(100, 2)) |
|||
self.domain_classifier.add_module('d_softmax', nn.LogSoftmax(dim=1)) |
|||
|
|||
def forward(self, input_data, alpha): |
|||
input_data = input_data.expand(input_data.data.shape[0], 1, 28, 28) |
|||
feature = self.feature(input_data) |
|||
feature = feature.view(-1, 50 * 4 * 4) |
|||
reverse_feature = ReverseLayerF.apply(feature, alpha) |
|||
class_output = self.class_classifier(feature) |
|||
domain_output = self.domain_classifier(reverse_feature) |
|||
|
|||
return class_output, domain_output |
|||
@ -0,0 +1,43 @@ |
|||
"""Params for DANN.""" |
|||
|
|||
import os |
|||
|
|||
# params for path |
|||
dataset_root = os.path.expanduser(os.path.join('~', 'Datasets')) |
|||
model_root = os.path.expanduser(os.path.join('~', 'Models', 'pytorch-DANN')) |
|||
|
|||
# params for datasets and data loader |
|||
dataset_mean_value = 0.5 |
|||
dataset_std_value = 0.5 |
|||
dataset_mean = (dataset_mean_value, dataset_mean_value, dataset_mean_value) |
|||
dataset_std = (dataset_std_value, dataset_std_value, dataset_std_value) |
|||
batch_size = 128 |
|||
image_size = 28 |
|||
|
|||
# params for source dataset |
|||
src_dataset = "SVHN" |
|||
src_model_trained = True |
|||
src_classifier_restore = os.path.join(model_root,src_dataset + '-source-classifier-final.pt') |
|||
|
|||
# params for target dataset |
|||
tgt_dataset = "MNIST" |
|||
tgt_model_trained = True |
|||
dann_restore = os.path.join(model_root , src_dataset + '-' + tgt_dataset + '-dann-final.pt') |
|||
|
|||
# params for pretrain |
|||
num_epochs_src = 100 |
|||
log_step_src = 10 |
|||
save_step_src = 20 |
|||
eval_step_src = 20 |
|||
|
|||
# params for training dann |
|||
num_epochs = 400 |
|||
log_step = 50 |
|||
save_step = 50 |
|||
eval_step = 20 |
|||
|
|||
manual_seed = 8888 |
|||
alpha = 0 |
|||
|
|||
# params for optimizing models |
|||
lr = 2e-4 |
|||
@ -0,0 +1,96 @@ |
|||
"""Utilities for ADDA.""" |
|||
|
|||
import os |
|||
import random |
|||
|
|||
import torch |
|||
import torch.backends.cudnn as cudnn |
|||
from torch.autograd import Variable |
|||
|
|||
import params |
|||
from datasets import get_mnist, get_mnistm, get_svhn |
|||
|
|||
|
|||
def make_variable(tensor, volatile=False): |
|||
"""Convert Tensor to Variable.""" |
|||
if torch.cuda.is_available(): |
|||
tensor = tensor.cuda() |
|||
return Variable(tensor, volatile=volatile) |
|||
|
|||
|
|||
def make_cuda(tensor): |
|||
"""Use CUDA if it's available.""" |
|||
if torch.cuda.is_available(): |
|||
tensor = tensor.cuda() |
|||
return tensor |
|||
|
|||
|
|||
def denormalize(x, std, mean): |
|||
"""Invert normalization, and then convert array into image.""" |
|||
out = x * std + mean |
|||
return out.clamp(0, 1) |
|||
|
|||
|
|||
def init_weights(layer): |
|||
"""Init weights for layers w.r.t. the original paper.""" |
|||
layer_name = layer.__class__.__name__ |
|||
if layer_name.find("Conv") != -1: |
|||
layer.weight.data.normal_(0.0, 0.02) |
|||
elif layer_name.find("BatchNorm") != -1: |
|||
layer.weight.data.normal_(1.0, 0.02) |
|||
layer.bias.data.fill_(0) |
|||
|
|||
|
|||
def init_random_seed(manual_seed): |
|||
"""Init random seed.""" |
|||
seed = None |
|||
if manual_seed is None: |
|||
seed = random.randint(1, 10000) |
|||
else: |
|||
seed = manual_seed |
|||
print("use random seed: {}".format(seed)) |
|||
random.seed(seed) |
|||
torch.manual_seed(seed) |
|||
if torch.cuda.is_available(): |
|||
torch.cuda.manual_seed_all(seed) |
|||
|
|||
|
|||
def get_data_loader(name, train=True): |
|||
"""Get data loader by name.""" |
|||
if name == "MNIST": |
|||
return get_mnist(train) |
|||
elif name == "MNISTM": |
|||
return get_mnistm(train) |
|||
elif name == "SVHN": |
|||
return get_svhn(train) |
|||
|
|||
|
|||
def init_model(net, restore): |
|||
"""Init models with cuda and weights.""" |
|||
# init weights of model |
|||
net.apply(init_weights) |
|||
|
|||
# restore model weights |
|||
if restore is not None and os.path.exists(restore): |
|||
net.load_state_dict(torch.load(restore)) |
|||
net.restored = True |
|||
print("Restore model from: {}".format(os.path.abspath(restore))) |
|||
else: |
|||
print("No trained model, train from scratch.") |
|||
|
|||
# check if cuda is available |
|||
if torch.cuda.is_available(): |
|||
cudnn.benchmark = True |
|||
net.cuda() |
|||
|
|||
return net |
|||
|
|||
|
|||
def save_model(net, filename): |
|||
"""Save trained model.""" |
|||
if not os.path.exists(params.model_root): |
|||
os.makedirs(params.model_root) |
|||
torch.save(net.state_dict(), |
|||
os.path.join(params.model_root, filename)) |
|||
print("save pretrained model to: {}".format(os.path.join(params.model_root, |
|||
filename))) |
|||
Loading…
Reference in new issue