def valid(args, dataloaders, models=None):
valid_source_dataloader = dataloaders['valid_source']
valid_target_dataloader = dataloaders['valid_target']
if models is None:
try:
load = torch.load(
f"./result/{args.source}2{args.target}/best_model.pth")
model = resnet(args)
#model = Model()
model.load_state_dict(load['M'])
except Exception as inst:
model = resnet(args)
model.eval()
model.to(args.device)
else:
model = models
model.eval()
src_accs = []
for i, (imgs, lbls) in enumerate(valid_source_dataloader):
bsize = imgs.size(0)
imgs, lbls = imgs.to(args.device), lbls.to(args.device)
output = model(imgs)
acc = accuracy_(output[:, :args.num_classes], lbls)
src_accs.append(acc)
print(f"\t [{i+1}/{len(valid_target_dataloader)}]", end=" \r")
tar_accs1 = []
tar_accs2 = []
for i, (imgs, lbls) in enumerate(valid_target_dataloader):
bsize = imgs.size(0)
imgs, lbls = imgs.to(args.device), lbls.to(args.device)
output = model(imgs)
acc = accuracy_(output[:, args.num_classes:], lbls)
tar_accs1.append(acc)
acc = accuracy_(
output[:, args.num_classes:] + output[:, :args.num_classes], lbls)
tar_accs2.append(acc)
print(f"\t [{i+1}/{len(valid_target_dataloader)}]", end=" \r")
mean_src_acc, mean_tar_acc1, mean_tar_acc2 = mean_(src_accs), mean_(
tar_accs1), mean_(tar_accs2)
#print(f"\t Valid, src acc:%.3f, tar acc:%.3f" % (mean_src_acc,
# mean_tar_acc))
return mean_src_acc, mean_tar_acc1, mean_tar_acc2
def valid(args, dataloaders, models=None):
valid_source_dataloader = dataloaders['valid_source']
valid_target_dataloader = dataloaders['valid_target']
if models is None:
load = torch.load(
f"./result/3_1/{args.source}2{args.target}/best_model.pth")
feature_extractor = FeatureExtractor()
feature_extractor.load_state_dict(load['F'])
feature_extractor.to(args.device)
label_predictor = LabelPredictor()
label_predictor.load_state_dict(load['C'])
label_predictor.to(args.device)
else:
feature_extractor = models['F']
label_predictor = models['C']
feature_extractor.eval()
label_predictor.eval()
src_accs = []
for i,(imgs, lbls) in enumerate(valid_source_dataloader):
bsize = imgs.size(0)
imgs, lbls = imgs.to(args.device), lbls.to(args.device)
features = feature_extractor(imgs)
class_output = label_predictor(features)
acc = accuracy_(class_output, lbls)
src_accs.append(acc)
print(f"\t [{i+1}/{len(valid_target_dataloader)}]", end=" \r")
tar_accs = []
for i,(imgs, lbls) in enumerate(valid_target_dataloader):
bsize = imgs.size(0)
imgs, lbls = imgs.to(args.device), lbls.to(args.device)
features = feature_extractor(imgs)
class_output = label_predictor(features)
acc = accuracy_(class_output, lbls)
tar_accs.append(acc)
print(f"\t [{i+1}/{len(valid_target_dataloader)}]", end=" \r")
mean_src_acc, mean_tar_acc = mean_(src_accs), mean_(tar_accs)
#print(f"\t Valid, src acc:%.3f, tar acc:%.3f" % (mean_src_acc,
# mean_tar_acc))
return mean_src_acc, mean_tar_acc
def train_base(args, base_loader, model, criterion, optimizer, epoch):
losses, accs = [Averager() for i in range(2)]
model.train()
for i, (imgs, lbls) in enumerate(base_loader):
imgs, lbls = imgs.to(args.device), lbls.to(args.device)
features = model(imgs)
predicts = model.classify(features)
loss = criterion(predicts, lbls)
acc = accuracy_(predicts, lbls)
optimizer.zero_grad()
loss.backward()
optimizer.step()
model.weight_norm()
losses.add(loss.item()), accs.add(acc)
print(" [%d/%d] base loss:%.3f, acc:%.2f" %
(i + 1, len(base_loader), losses.item(), accs.item()),
end=' \r')
wandb.log({
'base_train_loss': losses.item(),
'base_train_acc': accs.item()
})
print("Epoch %d, base loss:%.4f, acc:%.3f" %
(epoch, losses.item(), accs.item()))
def valid(args, valid_loader, model, distance, epoch):
model.eval()
if not hasattr(valid, "best_acc"):
valid.best_acc = 0
for i, batch in enumerate(valid_loader):
data, _ = [d.cuda() for d in batch]
p = args.shot * args.valid_way
data_support, data_query = data[:p], data[p:] #(shot*way) #(query*way)
# proto
proto = model(data_support).reshape(args.shot, args.valid_way, -1)
proto = proto.reshape(args.shot, args.valid_way, -1)
proto_hallu = model.hallucinate(proto.mean(dim=0), args.hallu_m)
proto_aug = torch.cat((proto, proto_hallu), dim=0) #(shot+m,way,z)
proto_aug = proto_aug.mean(dim=0) #(way,z)
# query
query = model(data_query) #(query*way,z)
query = query.view(1, args.query * args.valid_way, -1)
#query_hallu = model.hallucinate(query.squeeze(0), args.hallu_m)
#query_aug = torch.cat((query,query_hallu),dim=0) #((1+m),query*way,z)
#query_aug = query_aug.view((1+args.hallu_m)*args.query*args.valid_way, -1)
query_aug = query.squeeze(0)
# distance
logits = distance(query_aug, proto_aug) #(query*way, way)
label = torch.arange(args.valid_way).repeat(
#(1+args.hallu_m)*
args.query).long() #(query*way)
label = label.to(args.device)
loss = F.cross_entropy(logits, label)
acc = accuracy_(logits, label)
valid_loss.add(loss.item()), valid_acc.add(acc)
print(' [%d/%d], loss=%.3f acc=%.3f' %
(i + 1, len(valid_loader), valid_loss.item(), valid_acc.item()),
end=' \r')
wandb.log({'valid_loss': valid_loss.item(), 'valid_acc': valid_acc.item()})
print('Epoch %d, Meta-valid loss=%.4f acc=%.3f' %
(epoch, valid_loss.item(), valid_acc.item()))
if valid_acc.item() > valid.best_acc and epoch > args.epochs // 2:
os.system('mkdir -p checkpoints')
torch.save(model.state_dict(), f'checkpoints/{args.name}_best.pth')
print(' save weight')
valid.best_acc = valid_acc.item()
valid_loss.clean(), valid_acc.clean()
def train(args, train_loader, model, distance, optimizer, epoch):
model.train()
for i, batch in enumerate(train_loader):
data, _ = [d.to(args.device) for d in batch]
p = args.shot * args.train_way
data_support, data_query = data[:p], data[p:] #(shot*way) #(query*way)
# proto
proto = model(data_support)
proto = proto.reshape(args.shot, args.train_way, -1)
proto_hallu = model.hallucinate(proto.mean(dim=0), args.hallu_m)
proto_aug = torch.cat((proto, proto_hallu), dim=0) #(shot+m,way,z)
proto_aug = proto_aug.mean(dim=0) #(way,z)
# query
query = model(data_query) #(query*way,z)
query = query.view(1, args.query * args.train_way, -1)
query_hallu = model.hallucinate(query.squeeze(0), args.hallu_m)
query_aug = torch.cat((query, query_hallu),
dim=0) #((1+m),query*way,z)
query_aug = query_aug.view(
(1 + args.hallu_m) * args.query * args.train_way, -1)
# distance
logits = distance(query_aug, proto_aug) #((1+m)*query*way, way)
label = torch.arange(args.train_way).repeat(
(1 + args.hallu_m) * args.query).long() #(query*way)
label = label.to(args.device)
loss = F.cross_entropy(logits, label)
acc = accuracy_(logits, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss.add(loss.item()), train_acc.add(acc)
wandb.log({'train_loss': loss.item(), 'train_acc': acc})
print(' [%d/%d], loss=%.3f acc=%.3f' %
(i + 1, len(train_loader), train_loss.item(), train_acc.item()),
end=' \r')
print('Epoch %d, Meta-train loss=%.4f acc=%.3f' %
(epoch, train_loss.item(), train_acc.item()))
train_loss.clean(), train_acc.clean()
proto = proto.reshape(args.shot, args.train_way, -1)
proto_hallu = model.hallucinate(proto.mean(dim=0), args.hallu_m)
proto_aug = torch.cat((proto, proto_hallu), dim=0) #(shot+m,way,z)
proto_aug = proto_aug.mean(dim=0) #(way,z)
query = model(data_query) #(query*way,z)
logits = distance(query, proto_aug) #(query*way, way)
label = torch.arange(args.train_way).repeat(
args.query).long() #(query*way)
label = label.to(args.device)
loss = F.cross_entropy(logits, label)
acc = accuracy_(logits, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# adjust lr
lr_scheduler.step()
train_loss.add(loss.item()), train_acc.add(acc)
wandb.log({'train_loss': loss.item(), 'train_acc': acc})
print(' [%d/%d], loss=%.3f acc=%.3f' %
(i + 1, len(train_loader), train_loss.item(),
train_acc.item()),
end=' \r')
print('Epoch %d, train loss=%.4f acc=%.3f' %
def train(args, dataloaders):
train_source_dataloader = dataloaders['train_source']
train_target_dataloader = dataloaders['train_target']
# model
feature_extractor = FeatureExtractor().to(args.device)
label_predictor = LabelPredictor().to(args.device)
models = {'F':feature_extractor,
'C':label_predictor}
# loss
class_criterion = nn.CrossEntropyLoss()
# optimizer
optimizer = optim.Adam(list(label_predictor.parameters())+
list(feature_extractor.parameters()), lr=args.lr)
# lr schedule
steplr_after = StepLR(optimizer, step_size=10, gamma=1.0)
lr_scheduler = WarmupScheduler(optimizer, multiplier=1,
total_epoch=args.warmup_epochs,
after_scheduler=steplr_after)
#############################################################
########################### Train ###########################
train_source_dataloader_iter = circle_iterator(train_source_dataloader)
train_target_dataloader_iter = circle_iterator(train_target_dataloader)
best_tar_acc = 0
for epoch in range(1, args.epochs+1):
total_losses = []
class_losses = []
source_accs = []
len_dataloader = min(len(train_source_dataloader),
len(train_target_dataloader))
for i in range(len_dataloader):
src_imgs, src_lbls = next(train_source_dataloader_iter)
tar_imgs, tar_lbls = next(train_target_dataloader_iter)
src_imgs, src_lbls = src_imgs.to(args.device),\
src_lbls.to(args.device)
tar_imgs, tar_lbls = tar_imgs.to(args.device),\
tar_lbls.to(args.device)
######## Train Feature_extractor & Label_predictor #########
features = feature_extractor(src_imgs)
class_logits = label_predictor(features) #src
C_loss = class_criterion(class_logits, src_lbls)
loss = C_loss
loss.backward()
nn.utils.clip_grad_norm_(feature_extractor.parameters(), 0.001)
optimizer.step()
optimizer.zero_grad()
################### logging #####################
src_acc = accuracy_(class_logits, src_lbls)
class_losses.append(C_loss.item())
total_losses.append(loss.item())
source_accs.append(src_acc)
wandb.log({'class_loss': C_loss.item(),
'total_loss': loss.item(),
'source_acc': src_acc})
if i%5 == 0:
print(" [%d/%d] Loss Total:%.2f C:%.2f, src_acc:%.2f"
% ( i+1, len_dataloader,
loss.item(),
C_loss.item(),
src_acc), end=' \r')
# lr scheduler update
lr_scheduler.step()
# validating
valid_src_acc, valid_tar_acc = valid(args, dataloaders, models)
wandb.log({'valid_src_acc': valid_src_acc,
'valid_tar_acc': valid_tar_acc})
print(f" Epoch %d, Loss Total:%.3f C:%.3f"
", src_acc:%.3f tar_acc:%.3f " % (
epoch,
mean_(total_losses),
mean_(class_losses),
valid_src_acc,
valid_tar_acc))
if valid_tar_acc > 0.2 and best_tar_acc < valid_tar_acc:
best_tar_acc = valid_tar_acc
save_dir = f"./result/3_1/{args.source}2{args.target}/"
os.system(f"mkdir -p {save_dir}")
torch.save({'F':feature_extractor.state_dict(),
'C':label_predictor.state_dict()},
f"{save_dir}/best_model.pth")
print("\t save best weight")
def train(args, dataloaders):
train_source_dataloader = dataloaders['train_source']
train_target_dataloader = dataloaders['train_target']
# model
feature_extractor = FeatureExtractor().to(args.device)
#feature_extracter.apply(init_weights)
label_predictor = LabelPredictor().to(args.device)
domain_classifier = DomainClassifier().to(args.device)
models = {
'F': feature_extractor,
'C': label_predictor,
'D': domain_classifier
}
# loss
class_criterion = nn.CrossEntropyLoss()
domain_criterion = nn.BCEWithLogitsLoss()
# optimizer
optimizerF = optim.Adam(feature_extractor.parameters(), lr=args.lr)
optimizerC = optim.Adam(label_predictor.parameters(), lr=args.lr)
optimizerD = optim.Adam(domain_classifier.parameters(), lr=args.lr)
# lr schedule
'''
steplr_after = StepLR(optimizer, step_size=10, gamma=1.0)
lr_scheduler = WarmupScheduler(optimizer, multiplier=1,
total_epoch=args.warmup_epochs,
after_scheduler=steplr_after)
'''
#############################################################
########################### Train ###########################
train_source_dataloader_iter = circle_iterator(train_source_dataloader)
train_target_dataloader_iter = circle_iterator(train_target_dataloader)
best_tar_acc = 0
for epoch in range(1, args.epochs + 1):
total_losses = []
class_losses = []
domain_losses = []
source_accs = []
# calculate Alpha: how important domain loss is
#p = float(epoch) / args.epochs
#alpha = torch.tensor( 2. / (1.+np.exp(-10. * p)) -1 )
len_dataloader = min(len(train_source_dataloader),
len(train_target_dataloader))
for i in range(len_dataloader):
src_imgs, src_lbls = next(train_source_dataloader_iter)
tar_imgs, tar_lbls = next(train_target_dataloader_iter)
src_imgs, src_lbls = src_imgs.to(args.device),\
src_lbls.to(args.device)
tar_imgs, tar_lbls = tar_imgs.to(args.device),\
tar_lbls.to(args.device)
################### Train Domain_predictor ###################
mixed_imgs = torch.cat([src_imgs, tar_imgs], dim=0)
domain_label = torch.zeros(src_imgs.size(0) + tar_imgs.size(0),
1).to(args.device)
domain_label[:src_imgs.size(0)] = 1
features = feature_extractor(mixed_imgs)
domain_logits = domain_classifier(features.detach())
D_loss = domain_criterion(domain_logits, domain_label)
D_loss.backward()
optimizerD.step()
######## Train Feature_extractor & Label_predictor #########
class_logits = label_predictor(features[:src_imgs.size(0)]) #src
domain_logits = domain_classifier(features)
C_loss = class_criterion(class_logits, src_lbls)
D_loss = domain_criterion(domain_logits, domain_label)
loss = C_loss - args.lamb * D_loss
loss.backward()
nn.utils.clip_grad_norm_(feature_extractor.parameters(), 0.001)
optimizerC.step()
optimizerF.step()
optimizerF.zero_grad()
optimizerD.zero_grad()
optimizerC.zero_grad()
################### logging #####################
src_acc = accuracy_(class_logits, src_lbls)
class_losses.append(C_loss.item())
domain_losses.append(D_loss.item())
total_losses.append(loss.item())
source_accs.append(src_acc)
wandb.log({
'class_loss': C_loss.item(),
'domain_loss': D_loss.item(),
'total_loss': loss.item(),
'source_acc': src_acc
})
if i % 5 == 0:
print(" [%d/%d] Loss Total:%.2f C:%.2f D:%.2f, src_acc:%.2f" %
(i + 1, len_dataloader, loss.item(), C_loss.item(),
D_loss.item(), src_acc),
end=' \r')
# lr scheduler update
#lr_scheduler.step()
# validating
valid_src_acc, valid_tar_acc = valid(args, dataloaders, models)
wandb.log({
'valid_src_acc': valid_src_acc,
'valid_tar_acc': valid_tar_acc
})
print(f" Epoch %d, Loss Total:%.3f C:%.3f D:%.3f"
", src_acc:%.3f tar_acc:%.3f " %
(epoch, mean_(total_losses), mean_(class_losses),
mean_(domain_losses), valid_src_acc, valid_tar_acc))
if valid_tar_acc > 0.4 and best_tar_acc < valid_tar_acc:
best_tar_acc = valid_tar_acc
save_dir = f"./result/3_2/{args.source}2{args.target}/"
os.system(f"mkdir -p {save_dir}")
torch.save(
{
'F': feature_extractor.state_dict(),
'C': label_predictor.state_dict(),
'D': domain_classifier.state_dict()
}, f"{save_dir}/best_model.pth")
print("\t save best weight")
