import torch
import lib.modeling.resnet as resnet
import lib.modeling.semseg_heads as snet
import torch.nn as nn
import torch.optim as optim
import utils.resnet_weights_helper as resnet_utils
from torch.autograd import Variable
from roi_data.loader import RoiDataLoader, MinibatchSampler, collate_minibatch, collate_minibatch_semseg
from datasets.roidb import combined_roidb_for_training, combined_roidb_for_training_semseg
import os
import numpy as np
import nn as mynn
import cv2
from modeling.model_builder_3DSD import Generalized_3DSD
from modeling.model_builder_PSP3D import DispSeg
from core.config import cfg, cfg_from_file, cfg_from_list, assert_and_infer_cfg

#load net
class load_net(nn.Module):
    def __init__(self):
        super(load_net, self).__init__()

        build=snet.ModelBuilder()
        fc_dim = 2048
        self.encoder = build.build_encoder(
                arch= 'resnet50_dilated8',
                fc_dim=fc_dim)
        self.decoder = build.build_decoder(
                arch = 'ppm_bilinear',
                num_class=19,
                fc_dim=fc_dim,
                use_softmax=False)
    def _init_modules(self):
        resnet_utils.load_pretrained_imagenet_weights(self)

    def forward(self, data):
        pred=self.decoder(self.encoder(data, return_feature_maps=True))
        
        pred = nn.functional.interpolate(
                pred, size=[128,128],
                mode='bilinear', align_corners=False)
        pred = nn.functional.log_softmax(pred, dim=1)
        return pred
def dataloader(bs, gpus):
    inputs = {}
    inputs['data'] = Variable(torch.randn(2*bs, 3, 128, 128)).to('cuda')
    inputs['semseg_label_0'] = Variable(torch.LongTensor(
        np.random.randint(0, 19, (bs, 128//8, 128//8), dtype=np.long))).to('cuda')
    inputs['disp_label_0'] = Variable(torch.rand(bs,  128//8, 128//8)).to('cuda')
    inputs['disp_scans'] = Variable(torch.arange(0,
            cfg.DISP.MAX_DISPLACEMENT).float().view(1,cfg.DISP.MAX_DISPLACEMENT,1,1).repeat(bs,1,1,1)).to('cuda')
    inputs['semseg_scans'] = Variable(torch.arange(0,
            cfg.MODEL.NUM_CLASSES).float().view(1, cfg.MODEL.NUM_CLASSES, 1, 1).repeat(bs,1,1,1)).to('cuda')
    return inputs

cfg_file = 'e2e_segdisp-R-50_3Dpool_1x.yaml'
cfg_from_file(cfg_file)
print (cfg.SEM)
print (cfg.DISP)
#cfg_from_list(cfg_file)
#assert_and_infer_cfg()
devices_ids=[5]
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(ids) for ids in devices_ids])
torch.backends.cudnn.benchmark=True
#torch.cuda.set_device(3)
len_gpus = len(devices_ids)
batch_size = 2 * len_gpus
#net = mynn.DataParallel(load_net().to('cuda'), minibatch=True)
net = mynn.DataParallel(DispSeg().to('cuda'), minibatch=True)
optimizer = optim.SGD(net.parameters(), lr=0.000875, momentum=0.9)
criterion = nn.NLLLoss(ignore_index=255)
#dataloader= dataloader(batch_size, len_gpus)
for i in range(10):
#for i, inputs in zip(range(1000), dataloader):
    inputs = dataloader(batch_size, len_gpus)
    for key in inputs:
        inputs[key] = torch.chunk(inputs[key], chunks=len_gpus, dim=0)
    optimizer.zero_grad()
    loss=net(**inputs)
    optimizer.step()
    for k in loss['losses'].keys():
        print (loss['losses'][k].item())


[((2264, 2287), 'core.config.cfg_from_file', 'cfg_from_file', (['cfg_file'], {}), '(cfg_file)\n', (2277, 2287), False, 'from core.config import cfg, cfg_from_file, cfg_from_list, assert_and_infer_cfg\n'), ((2788, 2816), 'torch.nn.NLLLoss', 'nn.NLLLoss', ([], {'ignore_index': '(255)'}), '(ignore_index=255)\n', (2798, 2816), True, 'import torch.nn as nn\n'), ((786, 805), 'lib.modeling.semseg_heads.ModelBuilder', 'snet.ModelBuilder', ([], {}), '()\n', (803, 805), True, 'import lib.modeling.semseg_heads as snet\n'), ((1162, 1213), 'utils.resnet_weights_helper.load_pretrained_imagenet_weights', 'resnet_utils.load_pretrained_imagenet_weights', (['self'], {}), '(self)\n', (1207, 1213), True, 'import utils.resnet_weights_helper as resnet_utils\n'), ((1340, 1430), 'torch.nn.functional.interpolate', 'nn.functional.interpolate', (['pred'], {'size': '[128, 128]', 'mode': '"""bilinear"""', 'align_corners': '(False)'}), "(pred, size=[128, 128], mode='bilinear',\n    align_corners=False)\n", (1365, 1430), True, 'import torch.nn as nn\n'), ((1474, 1512), 'torch.nn.functional.log_softmax', 'nn.functional.log_softmax', (['pred'], {'dim': '(1)'}), '(pred, dim=1)\n', (1499, 1512), True, 'import torch.nn as nn\n'), ((3022, 3070), 'torch.chunk', 'torch.chunk', (['inputs[key]'], {'chunks': 'len_gpus', 'dim': '(0)'}), '(inputs[key], chunks=len_gpus, dim=0)\n', (3033, 3070), False, 'import torch\n'), ((2671, 2680), 'modeling.model_builder_PSP3D.DispSeg', 'DispSeg', ([], {}), '()\n', (2678, 2680), False, 'from modeling.model_builder_PSP3D import DispSeg\n'), ((1605, 1637), 'torch.randn', 'torch.randn', (['(2 * bs)', '(3)', '(128)', '(128)'], {}), '(2 * bs, 3, 128, 128)\n', (1616, 1637), False, 'import torch\n'), ((1827, 1861), 'torch.rand', 'torch.rand', (['bs', '(128 // 8)', '(128 // 8)'], {}), '(bs, 128 // 8, 128 // 8)\n', (1837, 1861), False, 'import torch\n'), ((1714, 1779), 'numpy.random.randint', 'np.random.randint', (['(0)', '(19)', '(bs, 128 // 8, 128 // 8)'], {'dtype': 'np.long'}), '(0, 19, (bs, 128 // 8, 128 // 8), dtype=np.long)\n', (1731, 1779), True, 'import numpy as np\n'), ((1907, 1949), 'torch.arange', 'torch.arange', (['(0)', 'cfg.DISP.MAX_DISPLACEMENT'], {}), '(0, cfg.DISP.MAX_DISPLACEMENT)\n', (1919, 1949), False, 'import torch\n'), ((2075, 2113), 'torch.arange', 'torch.arange', (['(0)', 'cfg.MODEL.NUM_CLASSES'], {}), '(0, cfg.MODEL.NUM_CLASSES)\n', (2087, 2113), False, 'import torch\n')]