def main(args: argparse.Namespace):
logger = CompleteLogger(args.log, args.phase)
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
cudnn.benchmark = True
# Data loading code
normalize = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if args.center_crop:
train_transform = T.Compose([
ResizeImage(256),
T.CenterCrop(224),
T.RandomHorizontalFlip(),
T.ToTensor(), normalize
])
else:
train_transform = T.Compose([
ResizeImage(256),
T.RandomResizedCrop(224),
T.RandomHorizontalFlip(),
T.ToTensor(), normalize
])
val_transform = T.Compose(
[ResizeImage(256),
T.CenterCrop(224),
T.ToTensor(), normalize])
dataset = datasets.__dict__[args.data]
train_source_dataset = dataset(root=args.root,
task=args.source,
download=True,
transform=train_transform)
train_source_loader = DataLoader(train_source_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
drop_last=True)
train_target_dataset = dataset(root=args.root,
task=args.target,
download=True,
transform=train_transform)
train_target_loader = DataLoader(train_target_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
drop_last=True)
val_dataset = dataset(root=args.root,
task=args.validation,
download=True,
transform=val_transform)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
if args.data == 'DomainNet':
test_dataset = dataset(root=args.root,
task=args.target,
split='test',
download=True,
transform=val_transform)
test_loader = DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
else:
test_loader = val_loader
train_source_iter = ForeverDataIterator(train_source_loader)
train_target_iter = ForeverDataIterator(train_target_loader)
# create model
print("=> using pre-trained model '{}'".format(args.arch))
G = models.__dict__[args.arch](pretrained=True).to(
device) # feature extractor
num_classes = train_source_dataset.num_classes
# two image classifier heads
F1 = ImageClassifierHead(G.out_features, num_classes,
args.bottleneck_dim).to(device)
F2 = ImageClassifierHead(G.out_features, num_classes,
args.bottleneck_dim).to(device)
# define optimizer
# the learning rate is fixed according to origin paper
optimizer_g = SGD(G.parameters(), lr=args.lr, weight_decay=0.0005)
optimizer_f = SGD([
{
"params": F1.parameters()
},
{
"params": F2.parameters()
},
],
momentum=0.9,
lr=args.lr,
weight_decay=0.0005)
# resume from the best checkpoint
if args.phase != 'train':
checkpoint = torch.load(logger.get_checkpoint_path('best'),
map_location='cpu')
G.load_state_dict(checkpoint['G'])
F1.load_state_dict(checkpoint['F1'])
F2.load_state_dict(checkpoint['F2'])
# analysis the model
if args.phase == 'analysis':
# extract features from both domains
feature_extractor = G.to(device)
source_feature = collect_feature(train_source_loader,
feature_extractor, device)
target_feature = collect_feature(train_target_loader,
feature_extractor, device)
# plot t-SNE
tSNE_filename = osp.join(logger.visualize_directory, 'TSNE.png')
tsne.visualize(source_feature, target_feature, tSNE_filename)
print("Saving t-SNE to", tSNE_filename)
# calculate A-distance, which is a measure for distribution discrepancy
A_distance = a_distance.calculate(source_feature, target_feature,
device)
print("A-distance =", A_distance)
return
if args.phase == 'test':
acc1 = validate(test_loader, G, F1, F2, args)
print(acc1)
return
# start training
best_acc1 = 0.
best_results = None
for epoch in range(args.epochs):
# train for one epoch
train(train_source_iter, train_target_iter, G, F1, F2, optimizer_g,
optimizer_f, epoch, args)
# evaluate on validation set
results = validate(val_loader, G, F1, F2, args)
# remember best acc@1 and save checkpoint
torch.save(
{
'G': G.state_dict(),
'F1': F1.state_dict(),
'F2': F2.state_dict()
}, logger.get_checkpoint_path('latest'))
if max(results) > best_acc1:
shutil.copy(logger.get_checkpoint_path('latest'),
logger.get_checkpoint_path('best'))
best_acc1 = max(results)
best_results = results
print("best_acc1 = {:3.1f}, results = {}".format(best_acc1, best_results))
# evaluate on test set
checkpoint = torch.load(logger.get_checkpoint_path('best'),
map_location='cpu')
G.load_state_dict(checkpoint['G'])
F1.load_state_dict(checkpoint['F1'])
F2.load_state_dict(checkpoint['F2'])
results = validate(test_loader, G, F1, F2, args)
print("test_acc1 = {:3.1f}".format(max(results)))
logger.close()
def main(args: argparse.Namespace):
logger = CompleteLogger(args.log, args.phase)
print(args)
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
cudnn.benchmark = True
# Data loading code
train_transform = utils.get_train_transform(
args.train_resizing,
random_horizontal_flip=not args.no_hflip,
random_color_jitter=False,
resize_size=args.resize_size,
norm_mean=args.norm_mean,
norm_std=args.norm_std)
val_transform = utils.get_val_transform(args.val_resizing,
resize_size=args.resize_size,
norm_mean=args.norm_mean,
norm_std=args.norm_std)
print("train_transform: ", train_transform)
print("val_transform: ", val_transform)
train_source_dataset, train_target_dataset, val_dataset, test_dataset, num_classes, args.class_names = \
utils.get_dataset(args.data, args.root, args.source, args.target, train_transform, val_transform)
train_source_loader = DataLoader(train_source_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
drop_last=True)
train_target_loader = DataLoader(train_target_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
drop_last=True)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
test_loader = DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
train_source_iter = ForeverDataIterator(train_source_loader)
train_target_iter = ForeverDataIterator(train_target_loader)
# create model
print("=> using model '{}'".format(args.arch))
G = utils.get_model(args.arch, pretrain=not args.scratch).to(
device) # feature extractor
# two image classifier heads
pool_layer = nn.Identity() if args.no_pool else None
F1 = ImageClassifierHead(G.out_features, num_classes, args.bottleneck_dim,
pool_layer).to(device)
F2 = ImageClassifierHead(G.out_features, num_classes, args.bottleneck_dim,
pool_layer).to(device)
# define optimizer
# the learning rate is fixed according to origin paper
optimizer_g = SGD(G.parameters(), lr=args.lr, weight_decay=0.0005)
optimizer_f = SGD([
{
"params": F1.parameters()
},
{
"params": F2.parameters()
},
],
momentum=0.9,
lr=args.lr,
weight_decay=0.0005)
# resume from the best checkpoint
if args.phase != 'train':
checkpoint = torch.load(logger.get_checkpoint_path('best'),
map_location='cpu')
G.load_state_dict(checkpoint['G'])
F1.load_state_dict(checkpoint['F1'])
F2.load_state_dict(checkpoint['F2'])
# analysis the model
if args.phase == 'analysis':
# extract features from both domains
feature_extractor = nn.Sequential(G, F1.pool_layer).to(device)
source_feature = collect_feature(train_source_loader,
feature_extractor, device)
target_feature = collect_feature(train_target_loader,
feature_extractor, device)
# plot t-SNE
tSNE_filename = osp.join(logger.visualize_directory, 'TSNE.pdf')
tsne.visualize(source_feature, target_feature, tSNE_filename)
print("Saving t-SNE to", tSNE_filename)
# calculate A-distance, which is a measure for distribution discrepancy
A_distance = a_distance.calculate(source_feature, target_feature,
device)
print("A-distance =", A_distance)
return
if args.phase == 'test':
acc1 = validate(test_loader, G, F1, F2, args)
print(acc1)
return
# start training
best_acc1 = 0.
best_results = None
for epoch in range(args.epochs):
# train for one epoch
train(train_source_iter, train_target_iter, G, F1, F2, optimizer_g,
optimizer_f, epoch, args)
# evaluate on validation set
results = validate(val_loader, G, F1, F2, args)
# remember best acc@1 and save checkpoint
torch.save(
{
'G': G.state_dict(),
'F1': F1.state_dict(),
'F2': F2.state_dict()
}, logger.get_checkpoint_path('latest'))
if max(results) > best_acc1:
shutil.copy(logger.get_checkpoint_path('latest'),
logger.get_checkpoint_path('best'))
best_acc1 = max(results)
best_results = results
print("best_acc1 = {:3.1f}, results = {}".format(best_acc1, best_results))
# evaluate on test set
checkpoint = torch.load(logger.get_checkpoint_path('best'),
map_location='cpu')
G.load_state_dict(checkpoint['G'])
F1.load_state_dict(checkpoint['F1'])
F2.load_state_dict(checkpoint['F2'])
results = validate(test_loader, G, F1, F2, args)
print("test_acc1 = {:3.1f}".format(max(results)))
logger.close()