def predict(): #nonlocal output_dict #model.eval() with torch.no_grad(): input_dict_cuda, label_dict_cuda = utils_data.nestedDictToDevice((input_dict, label_dict), device=device) output_dict_cuda = model(input_dict_cuda) output_dict = utils_data.nestedDictToDevice(output_dict_cuda, device='cpu')
def _inference(engine, batch): # now compute error model.eval() with torch.no_grad(): x, y = utils_data.nestedDictToDevice(batch, device=device) # make it work for dict input too y_pred = model(x) return y_pred, y
def _update(engine, batch): model.train() optimizer.zero_grad() x, y = utils_data.nestedDictToDevice(batch, device=device) # make it work for dict input too y_pred = model(x) loss = loss_fn(y_pred, y) loss.backward() optimizer.step() return loss.item(), y_pred
input_dict['bg_crop'] = input_dict['img_crop'] # white background input_dict['bg_crop'] = torch.from_numpy(np.ones((2,3,128,128))*3).float().cuda() input_dict['shuffled_appearance'] =torch.from_numpy(np.array([1,0]).astype('int')) # test on example # input_dict = ex input_dict['external_rotation_global'] = torch.from_numpy(np.eye(3)).float().cuda() #input_dict['external_rotation_global'] = torch.from_numpy(rotationMatrixXZY(theta=0, phi=0, psi=2.0)).float().cuda() input_dict['external_rotation_cam'] = torch.from_numpy(np.eye(3)).float().cuda() output_dict = None label_dict = None model.eval() with torch.no_grad(): input_dict_cuda = utils_data.nestedDictToDevice((input_dict), device=device) output_dict_cuda = model(input_dict_cuda) output_dict = utils_data.nestedDictToDevice(output_dict_cuda, device='cpu') skimage.io.imsave('whiteresults/' + filename + '_' + str(k) + '.png', tensor_to_img(output_dict['img_crop'][0])) #skimage.io.imsave('redin.png', X[i,:,:,:]) #skimage.io.imsave('results/' + str(i) + '.png', Y[i,:,:,:]) #pickle #pred = model.predict(X[:-2]) #print(Y[0]) #skimage.io.imsave('test1.png', Y[0]) #print(np.shape(Y[0])) #testimg = tensor_to_npimg(npimg_to_tensor(Y[0])) #skimage.io.imsave('test.png', testimg) #break #break