[ML] Knockoff Net: 랜덤한 입력으로 다른 모델을 모방할 수 있을까?

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1_conv = nn.Conv2d(1, 6, 5, 1) # 입력 1개, 출력 6개, 필터 크기는 5x5, 1칸 단위로 이동하면서 필터를 씌운다
        self.layer1_relu = nn.ReLU()             # 활성화 함수. ReLU(x) 는 max(x, 0)과 같다
        self.layer1_pool = nn.MaxPool2d(2)       # 각 2x2 칸마다 최대값 하나씩만을 남긴다
        self.layer2_conv = nn.Conv2d(6, 16, 5, 1)
        self.layer2_relu = nn.ReLU()
        self.layer2_pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(16*4*4, 10)
        self.softmax = nn.Softmax(dim=1)
    
    def forward(self, x):
        x1 = self.layer1_conv(x)  # 1x28x28 형식의 데이터가 6x24x24 형식으로 변환된다
        x2 = self.layer1_relu(x1) # 
        x3 = self.layer1_pool(x2) # 6x24x24 형식의 데이터가 6x12x12 형식으로 변환된다
        x4 = self.layer2_conv(x3) # 6x12x12 형식의 데이터가 16x8x8 형식으로 변환된다
        x5 = self.layer2_relu(x4) # 
        x6 = self.layer2_pool(x5) # 16x8x8 형식의 데이터가 16x4x4 형식으로 변환된다
        x7 = x6.view(-1, 256)     # 16x4x4 형식의 데이터가 256-벡터로 변환된다
        x8 = self.fc(x7)          # 256-벡터가 10-벡터로 변환된다
        x9 = self.softmax(x8)     # 최종 10-벡터의 각 값은 0~1 사이 값을 가지고 총합은 1이 된다
        return x9

net = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters())

tensor([2.6900e-10, 4.2002e-23, 6.6886e-01, 1.5307e-04, 1.0798e-23, 4.9928e-07,
         3.4142e-08, 1.7324e-13, 3.3099e-01, 4.0419e-12])

net_koff = MyModel() # 기존 net과 구조는 동일하다.
criterion_koff = nn.BCELoss() # 라벨이 아닌 두 '분포'에 대한 학습이므로 BCELoss 사용
optimizer_koff = optim.Adam(net_koff.parameters())

num_epoch = 10
st = time.time()
print(f"training with {len(train_data)} noises...")
for epoch in range(num_epoch):
    tot_loss = 0
    for x,y in train_loader:
        noise_x = torch.randn_like(x)    # train data와 동일한 크기의 랜덤 데이터를 생성
        y_target = net(noise_x).detach() # 모델 V(net)의 예측값을 타겟으로 설정
        y_pred   = net_koff(noise_x)     # 모델 A(net_koff)의 예측 y_pred가 y_target이 되도록 학습
        
        loss = criterion_koff(y_pred, y_target)
        tot_loss += loss
        optimizer_koff.zero_grad()
        loss.backward()
        optimizer_koff.step()
    print(f" - Epoch {epoch+1}/{num_epoch}. loss: {tot_loss}....(elapsed {time.time() - st}s)")
print(f"done (elapsed {time.time() - st}s)")

# 실제 MNIST의 test data로 test accuracy 측정
net_koff.eval()
acc, tot = 0, 0
for x, y in test_loader:
    y_pred = net_koff(x)
    acc += (y==y_pred.argmax(1)).sum()
    tot += len(y)
print(f"test accuracy of knockoff net: {acc}/{tot} ({100*acc/tot}%)")

class MyAnotherModel(nn.Module):
    def __init__(self):
        super(MyAnotherModel, self).__init__()
        self.layer1_conv = nn.Conv2d(1, 8, 5, 1)
        self.layer1_relu = nn.ReLU()
        self.layer1_pool = nn.MaxPool2d(2)
        self.layer2_conv = nn.Conv2d(8, 12, 5, 1)
        self.layer2_relu = nn.ReLU()
        self.layer2_pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(12*4*4, 50)
        self.fc_relu = nn.ReLU()
        self.fc2 = nn.Linear(50, 10)
        self.softmax = nn.Softmax(dim=1)
    
    def forward(self, x):
        x1 = self.layer1_conv(x)  # 1x28x28 형식의 데이터가 8x24x24 형식으로 변환된다
        x2 = self.layer1_relu(x1) # 
        x3 = self.layer1_pool(x2) # 8x24x24 형식의 데이터가 8x12x12 형식으로 변환된다
        x4 = self.layer2_conv(x3) # 8x12x12 형식의 데이터가 12x8x8 형식으로 변환된다
        x5 = self.layer2_relu(x4) # 
        x6 = self.layer2_pool(x5) # 12x8x8 형식의 데이터가 12x4x4 형식으로 변환된다
        x7 = x6.view(-1, 192)     # 12x4x4 형식의 데이터가 192-벡터로 변환된다
        x8 = self.fc1(x7)         # 192-벡터가 50-벡터로 변환된다
        x9 = self.fc_relu(x8)
        x10 = self.fc2(x9)        # 50-벡터가 10-벡터로 변환된다
        x11 = self.softmax(x10)
        return x11