-
Notifications
You must be signed in to change notification settings - Fork 140
/
adapt.py
131 lines (104 loc) · 4.47 KB
/
adapt.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
"""Adversarial adaptation to train target encoder."""
import os
import torch
import torch.optim as optim
from torch import nn
import params
from utils import make_variable
def train_tgt(src_encoder, tgt_encoder, critic,
src_data_loader, tgt_data_loader):
"""Train encoder for target domain."""
####################
# 1. setup network #
####################
# set train state for Dropout and BN layers
tgt_encoder.train()
critic.train()
# setup criterion and optimizer
criterion = nn.CrossEntropyLoss()
optimizer_tgt = optim.Adam(tgt_encoder.parameters(),
lr=params.c_learning_rate,
betas=(params.beta1, params.beta2))
optimizer_critic = optim.Adam(critic.parameters(),
lr=params.d_learning_rate,
betas=(params.beta1, params.beta2))
len_data_loader = min(len(src_data_loader), len(tgt_data_loader))
####################
# 2. train network #
####################
for epoch in range(params.num_epochs):
# zip source and target data pair
data_zip = enumerate(zip(src_data_loader, tgt_data_loader))
for step, ((images_src, _), (images_tgt, _)) in data_zip:
###########################
# 2.1 train discriminator #
###########################
# make images variable
images_src = make_variable(images_src)
images_tgt = make_variable(images_tgt)
# zero gradients for optimizer
optimizer_critic.zero_grad()
# extract and concat features
feat_src = src_encoder(images_src)
feat_tgt = tgt_encoder(images_tgt)
feat_concat = torch.cat((feat_src, feat_tgt), 0)
# predict on discriminator
pred_concat = critic(feat_concat.detach())
# prepare real and fake label
label_src = make_variable(torch.ones(feat_src.size(0)).long())
label_tgt = make_variable(torch.zeros(feat_tgt.size(0)).long())
label_concat = torch.cat((label_src, label_tgt), 0)
# compute loss for critic
loss_critic = criterion(pred_concat, label_concat)
loss_critic.backward()
# optimize critic
optimizer_critic.step()
pred_cls = torch.squeeze(pred_concat.max(1)[1])
acc = (pred_cls == label_concat).float().mean()
############################
# 2.2 train target encoder #
############################
# zero gradients for optimizer
optimizer_critic.zero_grad()
optimizer_tgt.zero_grad()
# extract and target features
feat_tgt = tgt_encoder(images_tgt)
# predict on discriminator
pred_tgt = critic(feat_tgt)
# prepare fake labels
label_tgt = make_variable(torch.ones(feat_tgt.size(0)).long())
# compute loss for target encoder
loss_tgt = criterion(pred_tgt, label_tgt)
loss_tgt.backward()
# optimize target encoder
optimizer_tgt.step()
#######################
# 2.3 print step info #
#######################
if ((step + 1) % params.log_step == 0):
print("Epoch [{}/{}] Step [{}/{}]:"
"d_loss={:.5f} g_loss={:.5f} acc={:.5f}"
.format(epoch + 1,
params.num_epochs,
step + 1,
len_data_loader,
loss_critic.data[0],
loss_tgt.data[0],
acc.data[0]))
#############################
# 2.4 save model parameters #
#############################
if ((epoch + 1) % params.save_step == 0):
torch.save(critic.state_dict(), os.path.join(
params.model_root,
"ADDA-critic-{}.pt".format(epoch + 1)))
torch.save(tgt_encoder.state_dict(), os.path.join(
params.model_root,
"ADDA-target-encoder-{}.pt".format(epoch + 1)))
torch.save(critic.state_dict(), os.path.join(
params.model_root,
"ADDA-critic-final.pt"))
torch.save(tgt_encoder.state_dict(), os.path.join(
params.model_root,
"ADDA-target-encoder-final.pt"))
return tgt_encoder