def __init__(self, dim_input=1024, dim_embed=128, dim_mu=20):
super(VAE, self).__init__()
self.model_encode = PointNetCls(num_points=dim_input, k=dim_embed)
self.model_decode = PointGen(num_points=dim_input, k=dim_mu)
self.fc_mu = nn.Linear(dim_embed, dim_mu)
self.fc_var = nn.Linear(dim_embed, dim_mu)
self.sig = nn.Sigmoid()
self.relu = nn.ReLU()
## Load victim model
############################################################
print('=====> Loading victim model from checkpoint...')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load('checkpoints/%s.pth' % args.model_path)
args.note = args.note + ' model note:' + checkpoint['args'].note
if args.data == 'modelnet40':
num_classes = 40
elif args.data == 'shapenetpart':
num_classes = 16
if args.model == 'pointnet':
model = PointNetCls(num_classes, args.feature_transform)
model = model.to(device)
elif args.model == 'pointnet2':
model = PointNet2ClsMsg(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'dgcnn':
model = DGCNN(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'pointcnn':
model = PointCNNCls(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'rscnn':
from models.rscnn import RSCNN
transforms=test_tfs)
train_loader = DataLoader(train_data,
num_workers=4,
batch_size=args.batch_size,
shuffle=True,
drop_last=True)
test_loader = DataLoader(test_data,
num_workers=4,
batch_size=args.batch_size,
shuffle=True,
drop_last=False)
########################################
## Intiate model
########################################
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
attack_model = PointNetCls(NUM_CLASSES, args.feature_transform).to(device)
attack_model = attack_model.to(device)
if len(args.resume) > 1:
print('=====> Loading from checkpoint...')
checkpoint = torch.load('checkpoints/%s.pth' % args.resume)
args = checkpoint['args']
torch.manual_seed(args.seed)
print("Random Seed: ", args.seed)
attack_model.load_state_dict(checkpoint['model_state_dict'])
print('Successfully resumed!')
else:
print("there's no checkpoint")
generator = Generator(up_ratio=1).to(device)
def load_models(mode, device, args):
"""
:param mode: "cls" or "disc"
:param args:
:return:
"""
if mode == "cls":
model = PointNetCls(k=40, feature_transform=False) # False
model = model.to(device)
try:
if args.checkpoint:
print('===============================')
print("Loading pretrained cls model ...")
print('===============================')
model.load_state_dict(torch.load(args.checkpoint))
except Exception as e:
print(e)
sys.exit(0)
elif mode == "seg":
model = PointNetSeg(NUM_SEG_CLASSES=50)
model = init_net(model, device, init_type=args.init_disc)
try:
if args.checkpoint:
print('===============================')
print("Loading pretrained cls model ...")
print('===============================')
model.load_state_dict(torch.load(args.checkpoint))
except Exception as e:
print(e)
sys.exit(0)
elif mode == "seg_regu":
model = PointNetSeg_regulization(NUM_SEG_CLASSES=50)
model = init_net(model, device, init_type=args.init_disc)
try:
if args.checkpoint:
print('===============================')
print("Loading pretrained cls model ...")
print('===============================')
model.load_state_dict(torch.load(args.checkpoint))
except Exception as e:
print(e)
sys.exit(0)
elif mode == "disc":
model = DeepConvDiscNet(input_dim=args.disc_indim, output_dim=1) # input_dim=args.input_pts, output_dim=1)
model = init_net(model, device, init_type=args.init_disc)
#if args.checkpoint_disc:
# print('===============================')
# print("Loading pretrained discriminator model ...")
# print('===============================')
# model.load_state_dict(torch.load(args.checkpoint_disc))
# model = init_net(model, device, init_type=args.init_disc)
elif mode == "disc_seg":
model = PointwiseDiscNet(input_pts=args.input_pts,input_dim=args.disc_indim)
model = init_net(model, device, init_type=args.init_disc)
elif mode == "disc_dual":
shared_disc = BaseDiscNet(input_pts=args.input_pts, input_dim=args.disc_indim, output_dim=256)
shapeDisc = ShapeDiscNet(shared_output_dim=256, num_shapes=16)
pointDisc = PointDiscNet(shared_output_dim=256, input_pts=args.input_pts)
# model_point = SharedPointDiscNet(
# input_pts=args.input_pts,
# input_dim=args.disc_indim,
# shared_output_dim=128)
# model_shape = SharedShapeDiscNet(
# input_pts=args.input_pts,
# input_dim=args.disc_indim,
# shared_output_dim=128,
# num_shapes=16,
# )
shared_disc = init_net(shared_disc, device, init_type=args.init_disc)
shapeDisc = init_net(shapeDisc, device, init_type=args.init_disc)
pointDisc = init_net(pointDisc, device, init_type=args.init_disc)
return shared_disc, shapeDisc, pointDisc
elif mode == "disc_stack":
model = StackDiscNet(input_pts = args.input_pts, input_dim=args.disc_indim, num_shapes=16)
model = init_net(model, device, init_type=args.init_disc)
else:
raise ValueError("Invalid mode {}!".format(mode))
return model
########################################
## Intiate model
########################################
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if args.data == 'modelnet40':
num_classes = 40
elif args.data == 'shapenetpart':
num_classes = 16
if args.model == 'pointnet':
model = PointNetCls(num_classes, args.feature_transform)
model = model.to(device)
elif args.model == 'pointnet2':
model = PointNet2ClsMsg(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'dgcnn':
model = DGCNN(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'pointcnn':
model = PointCNNCls(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'rscnn':
from models.rscnn import RSCNN ## use torch 0.4.1.post2
num_classes = 40
if args.model == 'dgcnn_kcutmix':
model = DGCNN(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr * 100,
momentum=0.9,
weight_decay=1e-4)
scheduler_c = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
250,
eta_min=1e-3)
else:
if args.model == 'pointnet_kcutmix':
model = PointNetCls(num_classes, args.feature_transform)
model = model.to(device)
elif args.model == 'pointnet2_kcutmix':
model = PointNet2ClsMsg(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'rscnn_kcutmix':
from models.rscnn import RSCNN
import models.rscnn_utils.pointnet2_utils as pointnet2_utils
model = RSCNN(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
optimizer = torch.optim.Adam(model.parameters(),
type=int,
default=16,
help='input batch size')
parser.add_argument('--local_rank',
default=-1,
type=int,
help='node rank for distributed training')
args = parser.parse_args()
dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)
cudnn.benchmark = True
num_classes = 40
if args.model == 'pointnet':
model = PointNetCls(num_classes, args.feature_transform)
#load pretrain model
state_dict = torch.load('model/p1.pth', map_location='cpu')
try:
model.load_state_dict(state_dict['model_state_dict'])
except RuntimeError:
state_dict = {k[7:]: v for k, v in state_dict.items()}
model.load_state_dict(state_dict)
model = DistributedDataParallel(model.cuda(), device_ids=[args.local_rank])
print('======> Loading data')
train_sampler = torch.utils.data.distributed.DistributedSampler(valid_ds)
test_loader = torch.utils.data.DataLoader(valid_ds,
batch_size=args.batch_size,
num_workers=4,
args.feature_transform, args.augment = bool(args.feature_transform), bool(
args.augment)
### Set random seed
args.seed = args.seed if args.seed > 0 else random.randint(1, 10000)
# dataset path
DATA_PATH = '/home/zjl/data/modelnet40_normal_resampled/'
########################################
## Intiate model
########################################
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_classes = 40
if args.model == 'pointnet_kcutmix':
model = PointNetCls(num_classes, args.feature_transform)
model = model.to(device)
elif args.model == 'pointnet2_kcutmix':
model = PointNet2ClsMsg(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'dgcnn_kcutmix':
model = DGCNN(num_classes)
model = model.to(device)
model = nn.DataParallel(model)
elif args.model == 'rscnn_kcutmix':
from models.rscnn import RSCNN
import models.rscnn_utils.pointnet2_utils as pointnet2_utils
import models.rscnn_utils.pytorch_utils as pt_utils
model = RSCNN(num_classes)