<gh_stars>0
import torch
import os
import random
import numpy as np
from torch import nn
from options import args_parser
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

## 初始化一些参数
args = args_parser()
random.seed(args.seed)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

## 生成器
class Generator(nn.Module):
    def __init__(self,nc,n_latent,n_class):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(n_latent+n_class,256,4,1,0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=False),

            nn.ConvTranspose2d(256,128,3,2,1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=False),

            nn.ConvTranspose2d(128,64,4,2,1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=False),

            nn.ConvTranspose2d(64,nc,4,2,1),
            nn.Tanh()
        )

    def forward(self, x, y):
        x = torch.cat([x, y], dim=1)
        x = self.main(x)
        return x

## 判别器
class Discriminator(nn.Module):
    def __init__(self,nc,n_class):
        super(Discriminator, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(nc+n_class,64,4,2,1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2,inplace=False),

            nn.Conv2d(64,128,4,2,1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2,inplace=False),

            nn.Conv2d(128,256,4,2,1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2,inplace=False),
        )
        self.D = nn.Sequential(
            nn.Linear(3 * 3 * 256, 10)
        )
        self.DO = nn.Sequential(
            nn.Linear(10, 1),
            nn.Sigmoid()
        )

    def forward(self, x, y):
        y = y.repeat(1, 1, *x.shape[2:4])
        x = torch.cat([x, y], dim=1)
        x = self.main(x).view(x.shape[0], -1)
        x1 = self.D(x)
        x2 = self.DO(x1)
        return x1,torch.squeeze(x2)

## 分类器
class Classifier (nn.Module):
    def __init__(self, nc,n_class):
        super(Classifier, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(nc,128,4,2,1),
            nn.ReLU(),
            nn.Dropout(0.2),

            nn.Conv2d(128,256,4,2,1),
            nn.ReLU(),
            nn.Dropout(0.2)
        )
        self.C = nn.Sequential(
            nn.Linear(7 * 7 * 256, n_class)
        )

    def forward(self, x):
        x = self.main(x).view(x.shape[0], -1)
        x1 = self.C(x)
        return x1

if __name__ == '__main__':
    train_data = datasets.MNIST(root='./mnist',train=True,transform=transforms.Compose([transforms.ToTensor()]))
    train_loader = DataLoader(train_data, batch_size=args.local_bs, shuffle=True)
    G = Generator(1,100,10).to(device)
    D = Discriminator(1,10).to(device)
    C = Classifier(1,10).to(device)

    optimizerG = torch.optim.Adam(G.parameters(), lr=args.local_lr)
    optimizerD = torch.optim.Adam(D.parameters(), lr=args.local_lr)
    optimizerC = torch.optim.Adam(C.parameters(), lr=args.local_lr)
    # optimizerC = torch.optim.SGD(C.parameters(), lr=0.01,momentum=0.9, weight_decay=1e-5)
    
    dis_criterion = nn.BCELoss(reduction='mean')
    aux_criterion = nn.CrossEntropyLoss()

    for epoch in range(args.epochs):
        total = 0
        correct =0
        for step, (x, y) in enumerate(train_loader):
            x = x.to(device)
            y = y.to(device)
            z = torch.randn((x.shape[0], 100, 1, 1)).to(device)

            ## y_real:batch_size*10*1*1,10的维度对应标签
            y_real_temp = torch.zeros([y.shape[0], 10]).to(device)
            y_real_temp[np.arange(0, len(x)), y] = 1
            y_real_temp.requires_grad = True
            y_real = y_real_temp.unsqueeze(2).unsqueeze(3)

            batch_size = x.size(0)
            real_label = torch.zeros(batch_size).fill_(1).cuda()
            fake_label = torch.zeros(batch_size).fill_(0).cuda()
            # ====================================================================
            ## 固定C、G
            D.train()
            for p in C.parameters():
                p.requires_grad = False;
            for p in G.parameters():
                p.requires_grad = False;
            for p in D.parameters():
                p.requires_grad = True;            
              

            ## 将生成图像+真实标签判定为假
            G_x = G(z, y_real).detach()
            D_G_x = D(G_x, y_real)[1]
            D_loss_1 = dis_criterion(D_G_x,fake_label)

            ## 将真实图像+真实标签判定为真
            D_x = D(x, y_real)[1]
            D_loss_2 = dis_criterion(D_x,real_label)

            D_loss = D_loss_1  + D_loss_2
            optimizerD.zero_grad()
            D_loss.backward()
            optimizerD.step()
            # ====================================================================
            ## 固定C、D
            G.train()
            for p in C.parameters():
                p.requires_grad = False;
            for p in G.parameters():
                p.requires_grad = True;
            for p in D.parameters():
                p.requires_grad = False;
            
            ## 欺骗判别器，即生成图像+真实标签判定为真
            G_x = G(z, y_real)
            D_G_x = D(G_x, y_real)[1]
            G_loss_1 = dis_criterion(D_G_x,real_label)

            ## 不让C对生成图片准确分类，生成hard sample
            C_G_x = C(G_x)
            G_loss_2 = aux_criterion(C_G_x,y)

            G_loss = G_loss_1 - G_loss_2
            optimizerG.zero_grad()
            G_loss.backward()
            optimizerG.step()
            # ====================================================================
            ## 固定G、D
            C.train()
            for p in C.parameters():
                p.requires_grad = True;
            for p in G.parameters():
                p.requires_grad = False;
            for p in D.parameters():
                p.requires_grad = False; 

            ## 将G生成图片正确分类
            G_x = G(z, y_real).detach()
            C_G_x = C(G_x)
            C_loss_1 = aux_criterion(C_G_x,y)

            ## 将真实图片正确分类
            C_x = C(x)
            C_loss_2 = aux_criterion(C_x,y)

            C_loss = C_loss_1 + C_loss_2

            optimizerC.zero_grad()
            C_loss.backward()
            optimizerC.step()
            
            C_x = C(x)
            _,pred = C_x.max(1)
            total += y.size(0)
            correct += (pred == y).sum().item()
            # ====================================================================
        print("epoch:%d------accuracy:%.2f"%(epoch,(100.0 * correct / total)))
        save_dirs = 'local/tgan/'
        if not os.path.exists(save_dirs):
            os.makedirs(save_dirs)
        torch.save(C, save_dirs+'C_epoch_%d.pth' %epoch)
        torch.save(D, save_dirs+'D_epoch_%d.pth' %epoch)
        torch.save(G, save_dirs+'G_epoch_%d.pth' %epoch)
import torch
import os
import random
import numpy as np
from torch import nn
from options import args_parser
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

## 初始化一些参数
args = args_parser()
random.seed(args.seed)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

## 生成器
class Generator(nn.Module):
    def __init__(self,nc,n_latent,n_class):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(n_latent+n_class,256,4,1,0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=False),

            nn.ConvTranspose2d(256,128,3,2,1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=False),

            nn.ConvTranspose2d(128,64,4,2,1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=False),

            nn.ConvTranspose2d(64,nc,4,2,1),
            nn.Tanh()
        )

    def forward(self, x, y):
        x = torch.cat([x, y], dim=1)
        x = self.main(x)
        return x

## 判别器
class Discriminator(nn.Module):
    def __init__(self,nc,n_class):
        super(Discriminator, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(nc+n_class,64,4,2,1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2,inplace=False),

            nn.Conv2d(64,128,4,2,1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2,inplace=False),

            nn.Conv2d(128,256,4,2,1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2,inplace=False),
        )
        self.D = nn.Sequential(
            nn.Linear(3 * 3 * 256, 10)
        )
        self.DO = nn.Sequential(
            nn.Linear(10, 1),
            nn.Sigmoid()
        )

    def forward(self, x, y):
        y = y.repeat(1, 1, *x.shape[2:4])
        x = torch.cat([x, y], dim=1)
        x = self.main(x).view(x.shape[0], -1)
        x1 = self.D(x)
        x2 = self.DO(x1)
        return x1,torch.squeeze(x2)

## 分类器
class Classifier (nn.Module):
    def __init__(self, nc,n_class):
        super(Classifier, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(nc,128,4,2,1),
            nn.ReLU(),
            nn.Dropout(0.2),

            nn.Conv2d(128,256,4,2,1),
            nn.ReLU(),
            nn.Dropout(0.2)
        )
        self.C = nn.Sequential(
            nn.Linear(7 * 7 * 256, n_class)
        )

    def forward(self, x):
        x = self.main(x).view(x.shape[0], -1)
        x1 = self.C(x)
        return x1

if __name__ == '__main__':
    train_data = datasets.MNIST(root='./mnist',train=True,transform=transforms.Compose([transforms.ToTensor()]))
    train_loader = DataLoader(train_data, batch_size=args.local_bs, shuffle=True)
    G = Generator(1,100,10).to(device)
    D = Discriminator(1,10).to(device)
    C = Classifier(1,10).to(device)

    optimizerG = torch.optim.Adam(G.parameters(), lr=args.local_lr)
    optimizerD = torch.optim.Adam(D.parameters(), lr=args.local_lr)
    optimizerC = torch.optim.Adam(C.parameters(), lr=args.local_lr)
    # optimizerC = torch.optim.SGD(C.parameters(), lr=0.01,momentum=0.9, weight_decay=1e-5)
    
    dis_criterion = nn.BCELoss(reduction='mean')
    aux_criterion = nn.CrossEntropyLoss()

    for epoch in range(args.epochs):
        total = 0
        correct =0
        for step, (x, y) in enumerate(train_loader):
            x = x.to(device)
            y = y.to(device)
            z = torch.randn((x.shape[0], 100, 1, 1)).to(device)

            ## y_real:batch_size*10*1*1,10的维度对应标签
            y_real_temp = torch.zeros([y.shape[0], 10]).to(device)
            y_real_temp[np.arange(0, len(x)), y] = 1
            y_real_temp.requires_grad = True
            y_real = y_real_temp.unsqueeze(2).unsqueeze(3)

            batch_size = x.size(0)
            real_label = torch.zeros(batch_size).fill_(1).cuda()
            fake_label = torch.zeros(batch_size).fill_(0).cuda()
            # ====================================================================
            ## 固定C、G
            D.train()
            for p in C.parameters():
                p.requires_grad = False;
            for p in G.parameters():
                p.requires_grad = False;
            for p in D.parameters():
                p.requires_grad = True;            
              

            ## 将生成图像+真实标签判定为假
            G_x = G(z, y_real).detach()
            D_G_x = D(G_x, y_real)[1]
            D_loss_1 = dis_criterion(D_G_x,fake_label)

            ## 将真实图像+真实标签判定为真
            D_x = D(x, y_real)[1]
            D_loss_2 = dis_criterion(D_x,real_label)

            D_loss = D_loss_1  + D_loss_2
            optimizerD.zero_grad()
            D_loss.backward()
            optimizerD.step()
            # ====================================================================
            ## 固定C、D
            G.train()
            for p in C.parameters():
                p.requires_grad = False;
            for p in G.parameters():
                p.requires_grad = True;
            for p in D.parameters():
                p.requires_grad = False;
            
            ## 欺骗判别器，即生成图像+真实标签判定为真
            G_x = G(z, y_real)
            D_G_x = D(G_x, y_real)[1]
            G_loss_1 = dis_criterion(D_G_x,real_label)

            ## 不让C对生成图片准确分类，生成hard sample
            C_G_x = C(G_x)
            G_loss_2 = aux_criterion(C_G_x,y)

            G_loss = G_loss_1 - G_loss_2
            optimizerG.zero_grad()
            G_loss.backward()
            optimizerG.step()
            # ====================================================================
            ## 固定G、D
            C.train()
            for p in C.parameters():
                p.requires_grad = True;
            for p in G.parameters():
                p.requires_grad = False;
            for p in D.parameters():
                p.requires_grad = False; 

            ## 将G生成图片正确分类
            G_x = G(z, y_real).detach()
            C_G_x = C(G_x)
            C_loss_1 = aux_criterion(C_G_x,y)

            ## 将真实图片正确分类
            C_x = C(x)
            C_loss_2 = aux_criterion(C_x,y)

            C_loss = C_loss_1 + C_loss_2

            optimizerC.zero_grad()
            C_loss.backward()
            optimizerC.step()
            
            C_x = C(x)
            _,pred = C_x.max(1)
            total += y.size(0)
            correct += (pred == y).sum().item()
            # ====================================================================
        print("epoch:%d------accuracy:%.2f"%(epoch,(100.0 * correct / total)))
        save_dirs = 'local/tgan/'
        if not os.path.exists(save_dirs):
            os.makedirs(save_dirs)
        torch.save(C, save_dirs+'C_epoch_%d.pth' %epoch)
        torch.save(D, save_dirs+'D_epoch_%d.pth' %epoch)
        torch.save(G, save_dirs+'G_epoch_%d.pth' %epoch)
## 初始化一些参数 ## 生成器 ## 判别器 ## 分类器 # optimizerC = torch.optim.SGD(C.parameters(), lr=0.01,momentum=0.9, weight_decay=1e-5) ## y_real:batch_size*10*1*1,10的维度对应标签 # ==================================================================== ## 固定C、G ## 将生成图像+真实标签判定为假 ## 将真实图像+真实标签判定为真 # ==================================================================== ## 固定C、D ## 欺骗判别器，即生成图像+真实标签判定为真 ## 不让C对生成图片准确分类，生成hard sample # ==================================================================== ## 固定G、D ## 将G生成图片正确分类 ## 将真实图片正确分类 # ====================================================================