fazil
/
pytorch-dann-resnet
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

update src_only flag

master
wogong 5 years ago
parent
62bc1b6445
commit
baa3c8b05d
2 changed files with 10 additions and 130 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 127
      core/train.py
  2. 11
      experiments/synsigns_gtsrb.py

127
core/train.py
View File

@ -87,122 +87,7 @@ def train_dann(model, params, src_data_loader, tgt_data_loader, tgt_data_loader_
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()
global_step += 1
# print step info
logger.add_scalar('src_loss_class', src_loss_class.item(), global_step)
logger.add_scalar('src_loss_domain', src_loss_domain.item(), global_step)
logger.add_scalar('tgt_loss_domain', tgt_loss_domain.item(), global_step)
logger.add_scalar('loss', loss.item(), global_step)
# if ((step + 1) % params.log_step == 0):
# print(
# "Epoch [{:4d}/{}] Step [{:2d}/{}]: src_loss_class={:.6f}, src_loss_domain={:.6f}, tgt_loss_domain={:.6f}, loss={:.6f}"
# .format(epoch + 1, params.num_epochs, step + 1, len_dataloader, src_loss_class.data.item(),
# src_loss_domain.data.item(), tgt_loss_domain.data.item(), loss.data.item()))
# eval model
if ((epoch + 1) % params.eval_step == 0):
tgt_test_loss, tgt_acc, tgt_acc_domain = test(model, tgt_data_loader, device, flag='target')
src_test_loss, src_acc, src_acc_domain = test(model, src_data_loader, device, flag='source')
logger.add_scalar('src_test_loss', src_test_loss, global_step)
logger.add_scalar('src_acc', src_acc, global_step)
logger.add_scalar('src_acc_domain', src_acc_domain, global_step)
logger.add_scalar('tgt_test_loss', tgt_test_loss, global_step)
logger.add_scalar('tgt_acc', tgt_acc, global_step)
logger.add_scalar('tgt_acc_domain', tgt_acc_domain, global_step)
# save model parameters
if ((epoch + 1) % params.save_step == 0):
save_model(model, params.model_root,
params.src_dataset + '-' + params.tgt_dataset + "-dann-{}.pt".format(epoch + 1))
# save final model
save_model(model, params.model_root, params.src_dataset + '-' + params.tgt_dataset + "-dann-final.pt")
return model
def train_src_only(model, params, src_data_loader, tgt_data_loader, tgt_data_loader_eval, device, logger):
"""Train dann."""
####################
# 1. setup network #
####################
# setup criterion and optimizer
if not params.finetune_flag:
print("training non-office task")
optimizer = optim.SGD(model.parameters(), lr=params.lr, momentum=params.momentum)
else:
print("training office task")
parameter_list = [{
"params": model.features.parameters(),
"lr": 0.001
}, {
"params": model.fc.parameters(),
"lr": 0.001
}, {
"params": model.bottleneck.parameters()
}, {
"params": model.classifier.parameters()
}, {
"params": model.discriminator.parameters()
}]
optimizer = optim.SGD(parameter_list, lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()
####################
# 2. train network #
####################
global_step = 0
for epoch in range(params.num_epochs):
# set train state for Dropout and BN layers
model.train()
# 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
if params.lr_adjust_flag == 'simple':
lr = adjust_learning_rate(optimizer, p)
else:
lr = adjust_learning_rate_office(optimizer, p)
logger.add_scalar('lr', lr, global_step)
# prepare domain label
size_src = len(images_src)
size_tgt = len(images_tgt)
label_src = torch.zeros(size_src).long().to(device) # source 0
label_tgt = torch.ones(size_tgt).long().to(device) # target 1
# make images variable
class_src = class_src.to(device)
images_src = images_src.to(device)
images_tgt = images_tgt.to(device)
# zero gradients for optimizer
optimizer.zero_grad()
# train on source domain
src_class_output, src_domain_output = model(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 = model(input_data=images_tgt, alpha=alpha)
tgt_loss_domain = criterion(tgt_domain_output, label_tgt)
if params.src_only_flag:
loss = src_loss_class
# optimize dann
@ -217,11 +102,11 @@ def train_src_only(model, params, src_data_loader, tgt_data_loader, tgt_data_loa
logger.add_scalar('tgt_loss_domain', tgt_loss_domain.item(), global_step)
logger.add_scalar('loss', loss.item(), global_step)
# if ((step + 1) % params.log_step == 0):
# print(
# "Epoch [{:4d}/{}] Step [{:2d}/{}]: src_loss_class={:.6f}, src_loss_domain={:.6f}, tgt_loss_domain={:.6f}, loss={:.6f}"
# .format(epoch + 1, params.num_epochs, step + 1, len_dataloader, src_loss_class.data.item(),
# src_loss_domain.data.item(), tgt_loss_domain.data.item(), loss.data.item()))
if ((step + 1) % params.log_step == 0):
print(
"Epoch [{:4d}/{}] Step [{:2d}/{}]: src_loss_class={:.6f}, src_loss_domain={:.6f}, tgt_loss_domain={:.6f}, loss={:.6f}"
.format(epoch + 1, params.num_epochs, step + 1, len_dataloader, src_loss_class.data.item(),
src_loss_domain.data.item(), tgt_loss_domain.data.item(), loss.data.item()))
# eval model
if ((epoch + 1) % params.eval_step == 0):

11
experiments/synsigns_gtsrb.py
View File

@ -19,6 +19,7 @@ class Config(object):
config = os.path.join(model_root, 'config.txt')
finetune_flag = False
lr_adjust_flag = 'simple'
src_only_flag = False
# params for datasets and data loader
batch_size = 128
@ -35,13 +36,7 @@ class Config(object):
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 = 50
eval_step_src = 20
# params for training dann
# params for GPU device
gpu_id = '0'
## for digit
@ -53,7 +48,7 @@ class Config(object):
manual_seed = None
alpha = 0
# params for optimizing models
# params for SGD optimizer
lr = 0.01
momentum = 0.9

Write Preview
Loading…
Cancel
Save