10, size=(int(args.wm_num[0]/args.wm_batchsize), args.wm_batchsize))
elif args.dataset == 'mnist':
np_labels = np.random.randint(
10, size=(int(args.wm_num[1]/args.wm_batchsize), args.wm_batchsize))
wm_labels = torch.from_numpy(np_labels).cuda()
#wm_labels = SpecifiedLabel()
best_real_acc, best_wm_acc, best_wm_input_acc = 0, 0, 0
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
train_loss, test_loss = [[], []], [[], []]
train_acc, test_acc = [[], []], [[], []]
# Model
print('==> Building model..')
if args.dataset == 'mnist':
Hidnet = UnetGenerator_mnist()
Disnet = DiscriminatorNet_mnist()
elif args.dataset == 'cifar10':
Hidnet = UnetGenerator()
Disnet = DiscriminatorNet()
#Dnnet = LeNet5()
#Dnnet = VGG('VGG19')
#Dnnet = model
#Dnnet = gcv.models.resnet50(pretrained=False)
Dnnet = ResNet34()
#Dnnet = PreActResNet18()
#Dnnet = GoogLeNet()
# net = DenseNet121()
# net = ResNeXt29_2x64d()
#Dnnet = MobileNet()
np_labels = np.random.randint(
10, size=(int(args.wm_num[0] / args.wm_batchsize), args.wm_batchsize))
elif args.dataset == 'mnist':
np_labels = np.random.randint(
10, size=(int(args.wm_num[1] / args.wm_batchsize), args.wm_batchsize))
wm_labels = torch.from_numpy(np_labels) # .cuda()
best_real_acc, best_wm_acc, best_wm_input_acc = 0, 0, 0
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
train_loss, test_loss = [[], []], [[], []]
train_acc, test_acc = [[], []], [[], []]
# Model
print('==> Building model..')
if args.dataset == 'mnist':
Hidnet = UnetGenerator_mnist()
Disnet = DiscriminatorNet_mnist()
elif args.dataset == 'cifar10':
Hidnet = UnetGenerator()
Disnet = DiscriminatorNet()
#Dnnet = LeNet5()
Dnnet = VGG('VGG19')
#Dnnet = ResNet101()
#Dnnet = PreActResNet18()
#Dnnet = GoogLeNet()
#Dnnet = MobileNetV2()
#Dnnet = DPN26()
# 多卡训练的GPU加速
# Hidnet = nn.DataParallel(Hidnet.cuda())