FSL-GNN代码解读

FSL-GNN代码解读

main.py（主函数）

1、加载数据集：

train_loader = generator.Generator(args.dataset_root, args, partition=\'train\', dataset=args.dataset)

2、初始化或加载模型：

enc_nn = models.load_model(\'enc_nn\', args, io)
metric_nn = models.load_model(\'metric_nn\', args, io)

if enc_nn is None or metric_nn is None:
	enc_nn, metric_nn = models.create_models(args=args)
    softmax_module = models.SoftmaxModule()

models.create_models(args=args) : in models.py

def create_models(args):
    print (args.dataset)
    if \'omniglot\' == args.dataset:
        enc_nn = EmbeddingOmniglot(args, 64)
    elif \'mini_imagenet\' == args.dataset:
        enc_nn = EmbeddingImagenet(args, 128)
    else:
        raise NameError(\'Dataset \' + args.dataset + \' not knows\')
    return enc_nn, MetricNN(args, emb_size=enc_nn.emb_size)

class EmbeddingOmniglot():				# 特征提取
class EmbeddingImagenet():				# 略

class MetricNN(nn.Module):
	if self.metric_network == \'gnn_iclr_nl\':……		# 正常的网络
	self.gnn_obj = gnn_iclr.GNN_nl()			# in gnn_iclr.py
	
	elif self.metric_network == \'gnn_iclr_active\':……	# 主动学习
	self.gnn_obj = gnn_iclr.GNN_active()# in gnn_iclr.py
	
class SoftmaxModule():		# 线性分类

class GNN_nl(nn.Module) & class GNN_active(nn.Module) : in gnn_iclr.py

class GNN_nl(nn.Module):		# 图网络主要部分
	class Wcompute(nn.Module)	# W邻接矩阵计算
    class Gconv(nn.Module)		# 组图
		def gmul(input)		# 更新图节点特征，W直接返回

3、训练

# 权重衰减
weight_decay = 1e-6

# 优化器
opt_enc_nn = optim.Adam(enc_nn.parameters(), lr=args.lr, weight_decay=weight_decay)
opt_metric_nn = optim.Adam(metric_nn.parameters(), lr=args.lr, weight_decay=weight_decay)

# 梯度置零，也就是把loss关于weight的导数变成0
opt_enc_nn.zero_grad()
opt_metric_nn.zero_grad()

# 训练
loss_d_metric = train_batch(
	model=[enc_nn, metric_nn, 
	softmax_module],
	data=[batch_x, label_x, batches_xi, labels_yi, oracles_yi, hidden_labels])

# 更新参数
opt_enc_nn.step()
opt_metric_nn.step()

# 自适应参数
adjust_learning_rate(optimizers=[opt_enc_nn, opt_metric_nn], lr=args.lr, iter=batch_idx)

# 显示训练中loss的更新
if batch_idx % args.log_interval == 0:
	display_str = \'Train Iter: {}\'.format(batch_idx)
	display_str += \'\tLoss_d_metric: {:.6f}\'.format(total_loss/counter)
	io.cprint(display_str)

# 测试
def test_one_shot(args, model, test_samples=5000, partition=\'test\') 定义于 test.py 中
val_acc_aux = test.test_one_shot	# 验证集上测试
test_acc_aux = test.test_one_shot	# 测试集上测试
test.test_one_shot(					# 训练集上测试
	args, 
	model=[enc_nn, metric_nn, softmax_module],
	test_samples=test_samples, 
	partition=\'train\')				

# 测试完毕，将模型设置回训练状态
enc_nn.train()
metric_nn.train()

# 若在验证集上的效果继续变好，则更新
if val_acc_aux is not None and val_acc_aux >= val_acc:

# 保存模型
torch.save(enc_nn, \'checkpoints/%s/models/enc_nn.t7\' % args.exp_name)
torch.save(metric_nn, \'checkpoints/%s/models/metric_nn.t7\' % args.exp_name)

# 全部训练完毕后进行测试
test.test_one_shot