import numpy as np
import matplotlib.pyplot as plt
from torch.autograd import Function
from sklearn.manifold import TSNE
import torch
import mnist
import mnistm
import itertools
import os
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
def optimizer_scheduler(optimizer, p):
"""
Adjust the learning rate of optimizer
:param optimizer: optimizer for updating parameters
:param p: a variable for adjusting learning rate
:return: optimizer
"""
for param_group in optimizer.param_groups:
param_group['lr'] = 0.01 / (1. + 10 * p) ** 0.75
return optimizer
def one_hot_embedding(labels, num_classes=10):
"""Embedding labels to one-hot form.
Args:
labels: (LongTensor) class labels, sized [N,].
num_classes: (int) number of classes.
Returns:
(tensor) encoded labels, sized [N, #classes].
"""
y = torch.eye(num_classes)
return y[labels]
def save_model(encoder, classifier, discriminator, training_mode, save_name):
print('Save models ...')
save_folder = 'trained_models'
if not os.path.exists(save_folder):
os.makedirs(save_folder)
torch.save(encoder.state_dict(), 'trained_models/encoder_' + str(training_mode) + '_' + str(save_name) + '.pt')
torch.save(classifier.state_dict(), 'trained_models/classifier_' + str(training_mode) + '_' + str(save_name) + '.pt')
if training_mode == 'dann':
torch.save(discriminator.state_dict(), 'trained_models/discriminator_' + str(training_mode) + '_' + str(save_name) + '.pt')
print('Model is saved !!!')
def plot_embedding(X, y, d, training_mode, save_name):
x_min, x_max = np.min(X, 0), np.max(X, 0)
X = (X - x_min) / (x_max - x_min)
y = list(itertools.chain.from_iterable(y))
y = np.asarray(y)
plt.figure(figsize=(10, 10))
for i in range(len(d)): # X.shape[0] : 1024
# plot colored number
if d[i] == 0:
colors = (0.0, 0.0, 1.0, 1.0)
else:
colors = (1.0, 0.0, 0.0, 1.0)
plt.text(X[i, 0], X[i, 1], str(y[i]),
color=colors,
fontdict={'weight': 'bold', 'size': 9})
plt.xticks([]), plt.yticks([])
if save_name is not None:
plt.title(save_name)
save_folder = 'saved_plot'
if not os.path.exists(save_folder):
os.makedirs(save_folder)
fig_name = 'saved_plot/' + str(training_mode) + '_' + str(save_name) + '.png'
plt.savefig(fig_name)
print('{} is saved'.format(fig_name))
def visualize(encoder, training_mode, save_name):
# Draw 512 samples in test_data
source_test_loader = mnist.mnist_test_loader
target_test_loader = mnistm.mnistm_test_loader
# Get source_test samples
source_label_list = []
source_img_list = []
for i, test_data in enumerate(source_test_loader):
if i >= 16: # to get only 512 samples
break
img, label = test_data
label = label.numpy()
img = img.cuda()
img = torch.cat((img, img, img), 1) # MNIST channel 1 -> 3
source_label_list.append(label)
source_img_list.append(img)
source_img_list = torch.stack(source_img_list)
source_img_list = source_img_list.view(-1, 3, 28, 28)
# Get target_test samples
target_label_list = []
target_img_list = []
for i, test_data in enumerate(target_test_loader):
if i >= 16:
break
img, label = test_data
label = label.numpy()
img = img.cuda()
target_label_list.append(label)
target_img_list.append(img)
target_img_list = torch.stack(target_img_list)
target_img_list = target_img_list.view(-1, 3, 28, 28)
# Stack source_list + target_list
combined_label_list = source_label_list
combined_label_list.extend(target_label_list)
combined_img_list = torch.cat((source_img_list, target_img_list), 0)
source_domain_list = torch.zeros(512).type(torch.LongTensor)
target_domain_list = torch.ones(512).type(torch.LongTensor)
combined_domain_list = torch.cat((source_domain_list, target_domain_list), 0).cuda()
print("Extract features to draw T-SNE plot...")
combined_feature = encoder(combined_img_list) # combined_feature : 1024,2352
tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=3000)
dann_tsne = tsne.fit_transform(combined_feature.detach().cpu().numpy())
print('Draw plot ...')
save_name = save_name + '_' + str(training_mode)
plot_embedding(dann_tsne, combined_label_list, combined_domain_list, training_mode, save_name)
def visualize_input():
source_test_loader = mnist.mnist_test_loader
target_test_loader = mnistm.mnistm_test_loader
# Get source_test samples
source_label_list = []
source_img_list = []
for i, test_data in enumerate(source_test_loader):
if i >= 16: # to get only 512 samples
break
img, label = test_data
label = label.numpy()
img = img.cuda()
img = torch.cat((img, img, img), 1) # MNIST channel 1 -> 3
source_label_list.append(label)
source_img_list.append(img)
source_img_list = torch.stack(source_img_list)
source_img_list = source_img_list.view(-1, 3, 28, 28)
# Get target_test samples
target_label_list = []
target_img_list = []
for i, test_data in enumerate(target_test_loader):
if i >= 16:
break
img, label = test_data
label = label.numpy()
img = img.cuda()
target_label_list.append(label)
target_img_list.append(img)
target_img_list = torch.stack(target_img_list)
target_img_list = target_img_list.view(-1, 3, 28, 28)
# Stack source_list + target_list
combined_label_list = source_label_list
combined_label_list.extend(target_label_list)
combined_img_list = torch.cat((source_img_list, target_img_list), 0)
source_domain_list = torch.zeros(512).type(torch.LongTensor)
target_domain_list = torch.ones(512).type(torch.LongTensor)
combined_domain_list = torch.cat((source_domain_list, target_domain_list), 0).cuda()
print("Extract features to draw T-SNE plot...")
combined_feature = combined_img_list # combined_feature : 1024,3,28,28
combined_feature = combined_feature.view(1024, -1) # flatten
# print(type(combined_feature), combined_feature.shape)
tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=3000)
dann_tsne = tsne.fit_transform(combined_feature.detach().cpu().numpy())
print('Draw plot ...')
save_name = 'input_tsne_plot'
plot_embedding(dann_tsne, combined_label_list, combined_domain_list, 'input', 'mnist_n_mnistM')
def get_free_gpu():
os.system('nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp')
memory_available = [int(x.split()[2]) for x in open('tmp', 'r').readlines()]
return np.argmax(memory_available)
一键复制
编辑
Web IDE
原始数据
按行查看
历史