import torch
import torchvision
from torchvision import datasets, models, transforms
from datasets.tusimple.config import CONFIG
from datasets.kitti.augmentations import *
from datasets.kitti.kitti_loader import kittiLoader
from datasets.tusimple.tusimple_loader import tusimpleLoader
from CAN_frontend import CAN
import matplotlib.pyplot as plt
from metrics import runningScore
device = torch.device("cpu")
# dataset = kittiLoader('/mnt/data/tejus/kitti_road/data_road/', split="train")
dataset = tusimpleLoader('/home/tejus/Downloads/train_set/', split="train")
trainloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=1,
pin_memory=True)
checkpoint_dir = '/home/tejus/lane-seg-experiments/Segmentation/CAN_logger/frontend_only/runs/2018-10-06_14-51-26/best_val_model.pkl'
checkpoint = torch.load(checkpoint_dir)
import matplotlib.pyplot as plt
# def save_img(img, img_path, is_image):
# if is_image:
# img[:,0,:,:] *= 0.229
# img[:,1,:,:] *= 0.224
# img[:,2,:,:] *= 0.225
import torch
import torchvision
from torchvision import datasets, models, transforms
from datasets.tusimple.config import CONFIG
from datasets.kitti.augmentations import *
from datasets.kitti.kitti_loader import kittiLoader
from datasets.tusimple.tusimple_loader import tusimpleLoader
from CAN import CAN
import matplotlib.pyplot as plt
from metrics import runningScore
use_device = "cpu"
device = torch.device(use_device)
dataset = tusimpleLoader('/mnt/data/tejus/train_set/', split="val")
checkpoint_dir = 'best_val_model.pkl'
if use_device == "cuda:0":
valloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=True)
checkpoint = torch.load(checkpoint_dir)
else:
valloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=False)
checkpoint = torch.load(checkpoint_dir, map_location=lambda storage, loc: storage)
def save_img(img, img_path, is_image):
if is_image:
img[:,0,:,:] *= 0.229
img[:,1,:,:] *= 0.224
b.fill_(0)
params['features.'+str(layer_idx)+'.weight'] = w
params['features.'+str(layer_idx)+'.weight'] = b
#torch.save(net.state_dict(), 'test_identity.wts')
layer_idx = 56
w = params['features.'+str(layer_idx)+'.weight']
w.fill_(0)
for i in range(w.shape[0]):
w[i,i,0,0] = 1
params['features.'+str(layer_idx)+'.weight'] = w
augmentations = Compose([RandomRotate(5), RandomHorizontallyFlip()])
train_dataset = tusimpleLoader('/home/tejus/Downloads/train_set/', split="train", augmentations=augmentations)
trainloader = torch.utils.data.DataLoader(train_dataset, batch_size=2, shuffle=True, num_workers=1, pin_memory=True)
val_dataset = tusimpleLoader('/home/tejus/Downloads/train_set/', split="val", augmentations=augmentations)
valloader = torch.utils.data.DataLoader(val_dataset, batch_size=2, shuffle=True, num_workers=1, pin_memory=True)
# dataset = kittiLoader('/mnt/data/tejus/kitti_road/data_road/', split="train", augmentations=augmentations)
# trainloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True, num_workers=1, pin_memory=True)
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=0.0001)
# optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, net.parameters()), lr = 0.0001, momentum=0.9) # 0.00001
# loss_fn = nn.BCEWithLogitsLoss()
loss_fn = nn.CrossEntropyLoss()
def update_lr(optimizer, epoch):
if epoch == 25:
net.train()
### freeze weights of frontend network
i = 0
for k, v in params.items():
v.requires_grad = False
i += 1
if i == 32:
break
###
augmentations = Compose([RandomRotate(5), RandomHorizontallyFlip()])
train_dataset = tusimpleLoader('/mnt/data/tejus/train_set/',
split="train",
augmentations=augmentations)
trainloader = torch.utils.data.DataLoader(train_dataset,
batch_size=TRAIN_BATCH,
shuffle=True,
num_workers=TRAIN_BATCH,
pin_memory=True)
val_dataset = tusimpleLoader('/mnt/data/tejus/train_set/',
split="val",
augmentations=None)
valloader = torch.utils.data.DataLoader(val_dataset,
batch_size=VAL_BATCH,
shuffle=True,
num_workers=VAL_BATCH,
pin_memory=True)
