def make_model():
model = FCN()
if torch.cuda.is_available():
model.to('cuda')
model = nn.DataParallel(model)
return model
def testFCN():
from gpu_util import log_mem
from fcn import FCN, device
bs = 1
image = torch.randn([bs, 3, 512, 512]).to(device)
rays_o = torch.randn([bs, 3, 512, 512]).to(device)
rays_d = torch.randn([bs, 3, 512, 512]).to(device)
model = FCN()
model.to(device)
out = model(image, rays_o, rays_d)
# Fire the engines
# Fire the engines
if args.clf == 'fcn':
print('Initializing FCN...')
model = FCN(args.input_size, args.output_size)
elif args.clf == 'svm':
print('Initializing SVM...')
model = SVM(args.input_size, args.output_size)
elif args.clf == 'resnet18':
print('Initializing ResNet18...')
model = resnet.resnet18(num_channels=args.num_channels,
num_classes=args.output_size)
model.load_state_dict(torch.load(init_path))
print('Load init: {}'.format(init_path))
model = nn.DataParallel(model.to(device), device_ids=args.device_id)
if 'sgd' in args.paradigm:
optim = optim.SGD(params=model.parameters(), lr=args.lr)
elif 'adam' in args.paradigm:
optim = optim.Adam(params=model.parameters(), lr=args.lr)
loss_fn = multiClassHingeLoss() if args.clf == 'svm' else F.nll_loss
loss_type = 'hinge' if args.clf == 'svm' else 'nll'
print('+' * 80)
best = 0
x_ax = []
acc_train = []