def main(argv):
data_path = None
input_image = None
output = None
weight_path = None
mode=5
drop_rate=0
lab=s.lab
classification=False
temp=.4
try:
opts, args = getopt.getopt(argv,"w:p:b:m:ld:ct:i:o:",["weight-path=", "datapath=",'model=','lab','drop-rate=','input=','output='])
except getopt.GetoptError as error:
print(error)
print( 'demo.py -w <path to weights file> -p <path to folder of images> OR -i <path to single image> -l <no argument. use if lab should be used>\
-d <amount of dropout used in model> -c <no argument. Use if model is classifier> -t <temperature for annealed mean> -o <output path for images>')
sys.exit(2)
print("opts", opts)
for opt, arg in opts:
if opt in ("-w", "--weight-path"):
weight_path = arg
elif opt in ("--datapath", "-p"):
data_path = arg
elif opt=='-m':
if arg in ('custom','0'):
mode = 0
elif arg in ('u','1','unet'):
mode = 1
elif arg in ('ende','2'):
mode = 2
elif arg in ('richzhang','classende','3'):
mode = 3
elif arg in ('colorunet','cu','4'):
mode = 4
elif arg in ('mu','5','middle'):
mode = 5
elif opt in ('-l','--lab'):
lab=True
elif opt in ("-d", "--drop-rate"):
drop_rate = float(arg)
elif opt =='-c':
classification=True
lab=True
elif opt=='-t':
temp=float(arg)
elif opt in ('-i','--input'):
input_image = arg
elif opt in ('-o','--output'):
output = arg
if data_path is None and input_image is None:
print('Please select an image or folder')
sys.exit()
trafo=transforms.Compose([transforms.Grayscale(3 if lab else 1), transforms.Resize((96,96))])
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if data_path is not None:
dataset = ImageDataset(data_path,lab=lab,pretrafo=trafo)
loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)
if input_image is not None:
img=trafo(Image.open(input_image))
if lab:
img=color.rgb2lab(np.asarray(img)/255)[...,:1]-np.array([50])[None,None,:]
loader = [(transforms.ToTensor()(img)[None,...].float(),input_image)]
classes=(150 if classification else 2) if lab else 3
#define model
UNet=None
zoom=False
if mode == 0:
UNet=model(col_channels=classes)
elif mode ==1:
UNet=unet(drop_rate=drop_rate,classes=classes)
elif mode ==2:
UNet=generator(drop_rate,classes)
elif mode == 3:
UNet=richzhang(drop_rate,classes)
zoom=True
elif mode == 4:
UNet=color_unet(True,drop_rate,classes)
elif mode == 5:
UNet = middle_unet(True,drop_rate,classes)
#load weights
try:
UNet.load_state_dict(torch.load(weight_path, map_location=device))
print("Loaded network weights from", weight_path)
except FileNotFoundError:
print("Did not find weight files.")
sys.exit()
outpath=None
UNet.to(device)
UNet.eval()
with torch.no_grad():
for i,(X,name) in enumerate(loader):
X=X.to(device)
unet_col=UNet(X)
col=show_colorization(unet_col,original=X,lab=lab,cl=classification,zoom=zoom,T=temp,return_img=output is not None)
if output:
try:
fp,f=os.path.split(name)
except TypeError:
fp,f=os.path.split(name[0])
n,e=f.split('.')
f='.'.join((n+'_color','png'))
outpath=output if os.path.isdir(output) or os.path.isdir(os.path.basename(output)) else fp
Image.fromarray(toInt(col[0])).save(os.path.join(outpath,f))
if output:
print('Finished colorization. Go to "%s" to see the colorized version(s) of the image(s)'%os.path.realpath(outpath))
def main(argv):
# setting argument defaults
mbsize = s.batch_size
report_freq = s.report_freq
weight_path = s.weights_path
weights_name = s.weights_name
lr = s.learning_rate
save_freq = s.save_freq
mode = 3
image_loss_weight = s.image_loss_weight
epochs = s.epochs
beta1, beta2 = s.betas
infinite_loop = s.infinite_loop
data_path = s.data_path
drop_rate = 0
lab = True
weighted_loss = True
weight_lambda = .25
load_list = s.load_list
help = 'train_classification.py -b <batch size> -e <amount of epochs to train. standard: infinite> -r <report frequency> -w <path to weights folder> \
-n <name> -s <save freq.> -l <learning rate> -p <path to data set> -d <dropout rate> -m <mode: differnet models> --beta1 <beta1 for adam>\
--beta2 <beta2 for adam> --lab <No argument. If used lab colorspace is cused> --weighted <No argument. If used *NO* class weights are used> \
--lambda <hyperparameter for class weights>'
try:
opts, args = getopt.getopt(argv, "he:b:r:w:l:s:n:p:d:i:m:", [
'epochs=', "mbsize=", "report-freq=", 'weight-path=', 'lr=',
'save-freq=', 'weight-name=', 'data_path=', 'drop_rate='
'beta1=', 'beta2=', 'lab', 'image-loss-weight=', 'weighted',
'mode=', 'lambda='
])
except getopt.GetoptError:
print(help)
sys.exit(2)
print("opts", opts)
for opt, arg in opts:
if opt == '-h':
print(help)
sys.exit()
elif opt in ("-b", "--mbsize"):
mbsize = int(arg)
elif opt in ("-e", "--epochs"):
epochs = int(arg)
infinite_loop = False
elif opt in ('-r', '--report-freq'):
report_freq = int(arg)
elif opt in ("-w", "--weight-path"):
weight_path = arg
elif opt in ("-n", "--weight-name"):
weights_name = arg
elif opt in ("-s", "--save-freq"):
save_freq = int(arg)
elif opt in ("-l", "--lr"):
lr = float(arg)
elif opt in ("-p", "--data_path"):
data_path = str(arg)
elif opt in ("-d", "--drop_rate"):
drop_rate = float(arg)
elif opt == '-m':
if arg in ('richzhang', '0', 'ende'):
mode = 0
elif arg in ('u', '1', 'unet'):
mode = 1
elif arg in ('color', '2', 'cu'):
mode = 2
elif arg in ('mu', '3', 'middle'):
mode = 3
elif opt == '--beta1':
beta1 = float(arg)
elif opt == '--beta2':
beta2 = float(arg)
elif opt == '--lab':
lab = True
elif opt == '--weighted':
weighted_loss = not weighted_loss
elif opt == '--load-list':
load_list = not load_list
elif opt == '--lambda':
weight_lambda = float(arg)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = None
in_size = 256
if 'cifar' in data_path:
in_size = 32
dataset = 0
elif 'places' in data_path:
in_size = 224
dataset = 1
elif 'stl' in data_path:
in_size = 96
dataset = 2
in_shape = (3, in_size, in_size)
#out_shape=(s.classes,32,32)
betas = (beta1, beta2)
weight_path_ending = os.path.join(weight_path, weights_name + '.pth')
loss_path_ending = os.path.join(weight_path,
weights_name + "_" + s.loss_name)
trainset = load_trainset(data_path,
lab=lab,
normalize=False,
load_list=load_list)
trainloader = torch.utils.data.DataLoader(
trainset,
batch_size=mbsize,
shuffle=True,
num_workers=2 if dataset in (0, 1) else 0)
print("NETWORK PATH:", weight_path_ending)
#define output channels of the model
classes = 150
#define model
if mode == 0:
classifier = generator(drop_rate, classes)
elif mode == 1:
classifier = unet(True, drop_rate, classes)
elif mode == 2:
classifier = color_unet(True, drop_rate, classes)
elif mode == 3:
classifier = middle_unet(True, drop_rate, classes)
#load weights
try:
classifier.load_state_dict(torch.load(weight_path_ending))
print("Loaded network weights from", weight_path)
except FileNotFoundError:
print("Initialize new weights for the generator.")
#sys.exit(2)
classifier.to(device)
#save the hyperparameters to a JSON-file for better oranization
model_description_path_ending = os.path.join(weight_path,
s.model_description_name)
# initialize model dict
try:
with open(model_description_path_ending, "r") as file:
model_dict = json.load(file)
except FileNotFoundError:
model_dict = {}
prev_epochs = 0
# save settings in dict if new weights are beeing initialized
if not weights_name in model_dict.keys():
model_dict[weights_name] = {
"loss_name":
loss_path_ending,
"epochs":
0,
"batch_size":
mbsize,
"lr":
lr,
"lab":
lab,
"betas":
betas,
"image_loss_weight":
image_loss_weight,
"weighted_loss":
weighted_loss,
"model":
'classification ' +
['richzhang', 'U-Net', 'color U-Net', 'middle U-Net'][mode]
}
else:
#load specified parameters from model_dict
params = model_dict[weights_name]
#mbsize=params['batch_size']
betas = params['betas']
#lr=params['lr']
lab = params['lab']
image_loss_weight = params['image_loss_weight']
weighted_loss = params['weighted_loss']
loss_path_ending = params['loss_name']
#memorize how many epochs already were trained if we continue training
prev_epochs = params['epochs'] + 1
#optimizer
optimizer = optim.Adam(classifier.parameters(), lr=lr, betas=betas)
class_weight_path = 'resources/class-weights.npy'
if weighted_loss:
weights = np.load(class_weight_path)
if dataset == 0:
class_weight_path = 'resources/cifar-lab-class-weights.pt'
weights = torch.load(class_weight_path).numpy()
elif dataset == 2:
if weight_lambda:
class_weight_path = 'resources/probdist_lab.pt'
prob_dict = torch.load(class_weight_path)
prob = np.array(list(prob_dict.values()))
weights = 1 / ((1 - weight_lambda) * prob / prob.sum() +
weight_lambda / classes)
else:
class_weight_path = 'resources/class-weights-lab150-stl.pt'
weights = torch.load(class_weight_path)
print('Class-weights loaded from ' + class_weight_path)
criterion = softCossEntropyLoss(
weights=weights,
device=device) if weighted_loss else softCossEntropyLoss(weights=None,
device=device)
loss_hist = []
soft_onehot = torch.load('resources/smooth_onehot150.pt',
map_location=device)
classifier.train()
# run over epochs
for e in (range(prev_epochs, prev_epochs +
epochs) if not infinite_loop else count(prev_epochs)):
g_running = 0
#load batches
for i, batch in enumerate(trainloader):
if dataset == 0: #cifar 10
(image, _) = batch
elif dataset in (1, 2): #places
image = batch
#batch_size=image.shape[0]
if dataset == 0: #cifar/stl 10
image = np.transpose(image, (0, 2, 3, 1))
image = np.transpose(color.rgb2lab(image), (0, 3, 1, 2))
image = torch.from_numpy(
(image -
np.array([50, 0, 0])[None, :, None, None])).float()
X = image[:, :1, :, :].to(
device) #set X to the Lightness of the image
image = image[:, 1:, :, :].to(device) #image is a and b channel
#----------------------------------------------------------------------------------------
################################### Model optimization ##################################
#----------------------------------------------------------------------------------------
#clear gradients
optimizer.zero_grad()
#softmax activated distribution
model_out = classifier(X).double()
#create bin coded verion of ab ground truth
binab = ab2bins(image.transpose(1, 3).transpose(1, 2))
if mode == 0:
binab = F.interpolate(binab.float(),
scale_factor=(.25, .25)).long()
binab = torch.squeeze(binab, 1)
binab = soft_onehot[:, binab].transpose(0, 1).double()
#calculate loss
loss = criterion(model_out, binab).mean(0)
loss.backward()
optimizer.step()
g_running += loss.item()
loss_hist.append([e, loss.item()])
#report running loss
if (i + len(trainloader) * e) % report_freq == report_freq - 1:
print('Epoch %i, batch %i: \tloss=%.2e' %
(e + 1, i + 1, g_running / report_freq))
g_running = 0
if s.save_weights and (
i + len(trainloader) * e) % save_freq == save_freq - 1:
#save parameters
try:
torch.save(classifier.state_dict(), weight_path_ending)
#torch.save(crit.state_dict(),crit_path)
except FileNotFoundError:
os.makedirs(weight_path)
torch.save(classifier.state_dict(), weight_path_ending)
#torch.save(crit.state_dict(),crit_path)
print("Parameters saved")
if s.save_loss:
#save loss history to file
try:
f = open(loss_path_ending, 'a')
np.savetxt(f, loss_hist, '%e')
f.close()
except FileNotFoundError:
os.makedirs(s.loss_path)
np.savetxt(loss_path_ending, loss_hist, '%e')
loss_hist = []
#update epoch count in dict after each epoch
model_dict[weights_name]["epochs"] = e
#save it to file
try:
with open(model_description_path_ending, "w") as file:
json.dump(model_dict, file, sort_keys=True, indent=4)
except:
print('Could not save to model dictionary (JSON-file)')
def main(argv):
data_path = s.data_path
weight_path = s.weights_path
mode = 1
drop_rate = 0
lab = s.lab
classification = False
temp = .4
try:
opts, args = getopt.getopt(argv, "h:w:p:b:m:ld:ct:", [
"help", "weight-path=", "datapath=", 'model=', 'lab', 'drop-rate='
])
except getopt.GetoptError as error:
print(error)
print(
'test.py -w <path to weights file> -p <path to dataset> -l <no argument. use if lab should be used> -m <mode: different models>\
-d <amount of dropout used in model> -c <no argument. Use if model is classifier> -t <temperature for annealed mean> -o <output path for images>'
)
sys.exit(2)
print("opts", opts)
for opt, arg in opts:
if opt == '-h':
print('test.py -i <Boolean> -s <Boolean>')
sys.exit()
elif opt in ("-w", "--weight-path"):
weight_path = arg
elif opt in ("--datapath", "-p"):
data_path = arg
elif opt in ("--batchnorm", "-b"):
batch_norm = arg in ["True", "true", "1"]
elif opt == '-m':
if arg in ('custom', '0'):
mode = 0
elif arg in ('u', '1', 'unet'):
mode = 1
elif arg in ('ende', '2'):
mode = 2
elif arg in ('richzhang', 'classende', '3'):
mode = 3
elif arg in ('colorunet', 'cu', '4'):
mode = 4
elif arg in ('mu', '5', 'middle'):
mode = 5
elif opt in ('-l', '--lab'):
lab = True
elif opt in ("-d", "--drop-rate"):
drop_rate = float(arg)
elif opt == '-c':
classification = True
lab = True
elif opt == '-t':
temp = float(arg)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = None
if data_path == './cifar-10':
in_size = 32
dataset = 0
elif 'places' in data_path:
in_size = 224
dataset = 1
elif 'stl' in data_path:
in_size = 96
dataset = 2
in_shape = (3, in_size, in_size)
#out_shape=(s.classes,32,32)
trainset = load_trainset(data_path, train=False, lab=lab)
trainloader = torch.utils.data.DataLoader(
trainset,
batch_size=3,
shuffle=True,
num_workers=2 if dataset in (0, 1) else 0)
print("Loaded dataset from", data_path)
classes = (150 if classification else 2) if lab else 3
#define model
UNet = None
zoom = False
if mode == 0:
UNet = model(col_channels=classes)
elif mode == 1:
UNet = unet(drop_rate=drop_rate, classes=classes)
elif mode == 2:
UNet = generator(drop_rate, classes)
elif mode == 3:
UNet = richzhang(drop_rate, classes)
zoom = True
elif mode == 4:
UNet = color_unet(True, drop_rate, classes)
elif mode == 5:
UNet = middle_unet(True, drop_rate, classes)
#load weights
try:
UNet.load_state_dict(torch.load(weight_path, map_location=device))
print("Loaded network weights from", weight_path)
except FileNotFoundError:
print("Did not find weight files.")
#sys.exit(2)
UNet.to(device)
UNet.eval()
gray = torch.tensor([0.2989, 0.5870, 0.1140])[:, None, None].float()
with torch.no_grad():
for i, batch in enumerate(trainloader):
if dataset == 0: #cifar 10
(image, _) = batch
elif dataset in (1, 2): #places
image = batch
X = None
if lab:
if dataset == 0: #cifar 10
image = np.transpose(image, (0, 2, 3, 1))
image = np.transpose(color.rgb2lab(image), (0, 3, 1, 2))
image = torch.from_numpy(
(image -
np.array([50, 0, 0])[None, :, None, None])).float()
X = torch.unsqueeze(image[:, 0, :, :], 1).to(
device) #set X to the Lightness of the image
image = image[:, 1:, :, :] #image is a and b channel
else:
#convert to grayscale image
#using the matlab formula: 0.2989 * R + 0.5870 * G + 0.1140 * B and load data to gpu
X = (image.clone() * gray).sum(1).to(device).view(
-1, 1, *in_shape[1:])
image = image.float()
#print(X.min(),X.max())
#generate colorized version with unet
#for arr in (image[:,0,...],image[:,1,...],X):
# print(arr.min(),arr.max())
try:
unet_col = UNet(X)
except:
unet_col = UNet(torch.stack((X, X, X), 1)[:, :, 0, :, :])
#for arr in (unet_col[0,...],unet_col[1,...]):
# print(arr.min().item(),arr.max().item())
show_colorization(unet_col,
image,
X,
lab=lab,
cl=classification,
zoom=zoom,
T=temp)
def main(argv):
# setting argument defaults
mbsize = s.batch_size
report_freq=s.report_freq
weight_path=s.weights_path
weights_name=s.weights_name
lr=s.learning_rate
save_freq = s.save_freq
mode=3
image_loss_weight=s.image_loss_weight
epochs = s.epochs
beta1,beta2=s.betas
infinite_loop=s.infinite_loop
data_path = s.data_path
drop_rate = 0
lab = s.lab
load_list = s.load_list
help='train_gan.py -b <batch size> -e <amount of epochs to train. standard: infinite> -r <report frequency> -w <path to weights folder> \
-n <name> -s <save freq.> -l <learning rate> -p <path to data set> -d <dropout rate> -m <mode: differnet models> --beta1 <beta1 for adam>\
--beta2 <beta2 for adam> --lab <No argument. If used lab colorspace is used> --weighted <No argument. If used *NO* class weights are used> \
--lambda <hyperparameter for class weights>'
try:
opts, args = getopt.getopt(argv,"he:b:r:w:l:s:n:m:p:d:i:",
['epochs=',"mbsize=","report-freq=",'weight-path=', 'lr=','save-freq=','weight-name=','mode=','data_path=','drop_rate='
'beta1=','beta2=','lab','image-loss-weight=','load-list'])
except getopt.GetoptError:
print(help)
sys.exit(2)
print("opts" ,opts)
for opt, arg in opts:
if opt == '-h':
print(help)
sys.exit()
elif opt in ("-b", "--mbsize"):
mbsize = int(arg)
#elif opt in ("-p", "--data-path"):
# data_path = arg
elif opt in ("-e", "--epochs"):
epochs = int(arg)
infinite_loop=False
elif opt in ('-r','--report-freq'):
report_freq = int(arg)
elif opt in ("-w", "--weight-path"):
weight_path = arg
elif opt in ("-n", "--weight-name"):
weights_name = arg
elif opt in ("-s", "--save-freq"):
save_freq=int(arg)
elif opt in ("-l", "--lr"):
lr = float(arg)
elif opt=='-m':
if arg in ('custom','0'):
mode = 0
elif arg in ('u','1','unet'):
mode = 1
elif arg in ('ende','2'):
mode = 2
elif arg in ('mu','3','middle'):
mode = 3
elif arg in ('cu','4','middle'):
mode = 4
elif opt in ("-p", "--data_path"):
data_path = str(arg)
elif opt in ("-d", "--drop_rate"):
drop_rate = float(arg)
elif opt=='--beta1':
beta1 = float(arg)
elif opt=='--beta2':
beta2 = float(arg)
elif opt=='--lab':
lab=True
elif opt in ('-i','--image-loss-weight'):
image_loss_weight=float(arg)
elif opt in ('--load-list'):
load_list=True
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset=None
if 'cifar' in data_path:
in_size = 32
dataset = 0
elif 'places' in data_path:
in_size = 224
dataset = 1
elif 'stl' in data_path:
in_size = 96
dataset = 2
in_shape=(3,in_size,in_size)
betas=(beta1,beta2)
weight_path_ending=os.path.join(weight_path,weights_name+'.pth')
loss_path_ending = os.path.join(weight_path, weights_name + "_" + s.loss_name)
trainset = load_trainset(data_path,lab=lab,load_list=load_list)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=mbsize,
shuffle=True, num_workers=2)
print("NETWORK PATH:", weight_path_ending)
#define output channels of the model
classes=2 if lab else 3
#define model
UNet=None
try:
if mode ==0:
UNet=model(col_channels=classes)
elif mode ==1:
UNet=unet(drop_rate=drop_rate,classes=classes)
elif mode ==2:
UNet=generator(drop_rate,classes)
elif mode ==3:
UNet=middle_unet(drop_rate=drop_rate,classes=classes)
elif mode ==4:
UNet=color_unet(drop_rate=drop_rate,classes=classes)
#load weights
try:
UNet.load_state_dict(torch.load(weight_path_ending))
print("Loaded network weights from", weight_path)
except FileNotFoundError:
print("Initialize new weights for the generator.")
except RuntimeError:
#if the wrong mode was chosen: try the other one
UNet=model(col_channels=classes) if mode==1 else unet(classes=classes)
#load weights
try:
UNet.load_state_dict(torch.load(weight_path_ending))
print("Loaded network weights from", weight_path)
#change mode to the correct one
mode = (mode +1) %2
except FileNotFoundError:
print("Initialize new weights for the generator.")
UNet.to(device)
#save the hyperparameters to a JSON-file for better oranization
model_description_path_ending = os.path.join(weight_path, s.model_description_name)
# initialize model dict
try:
with open(model_description_path_ending, "r") as file:
model_dict = json.load(file)
except FileNotFoundError:
model_dict = {}
prev_epochs=0
# save settings in dict if new weights are beeing initialized
if not weights_name in model_dict.keys():
model_dict[weights_name] = {
"loss_name": loss_path_ending,
"epochs": 0,
"batch_size": mbsize,
"lr": lr,
"lab":lab,
"betas": betas,
"image_loss_weight": image_loss_weight,
"model":['custom','unet','encoder-decoder','middle-unet','color-unet'][mode]
}
else:
#load specified parameters from model_dict
params=model_dict[weights_name]
mbsize=params['batch_size']
betas=params['betas']
lr=params['lr']
lab=params['lab']
image_loss_weight=params['image_loss_weight']
loss_path_ending=params['loss_name']
#memorize how many epochs already were trained if we continue training
prev_epochs=params['epochs']+1
#define critic
crit=critic(in_size,classes=classes).to(device)
#load discriminator weights
crit_path=os.path.join(weight_path,weights_name+'_crit.pth')
try:
crit.load_state_dict(torch.load(crit_path))
print('Loaded weights for discriminator from %s'%crit_path)
except FileNotFoundError:
print('Initialize new weights for discriminator')
crit.apply(weights_init_normal)
#optimizer
optimizer_g=optim.Adam(UNet.parameters(),lr=lr,betas=betas)
optimizer_c=optim.Adam(crit.parameters(),lr=lr,betas=betas)
criterion = nn.BCELoss().to(device)
#additional gan loss: l1 loss
l1loss = nn.L1Loss().to(device)
loss_hist=[]
UNet.train()
crit.train()
gray = torch.tensor([0.2989 ,0.5870, 0.1140 ])[:,None,None].float()
ones = torch.ones(mbsize,device=device)
zeros= torch.zeros(mbsize,device=device)
# run over epochs
for e in (range(prev_epochs, prev_epochs + epochs) if not infinite_loop else count(prev_epochs)):
g_running,c_running=0,0
#load batches
for i,batch in enumerate(trainloader):
if dataset == 0: #cifar 10
(image,_) = batch
elif dataset in (1,2): #places
image = batch
batch_size=image.shape[0]
X=None
#differentiate between the two available color spaces RGB and Lab
if lab:
if dataset == 0: #cifar 10
image=np.transpose(image,(0,2,3,1))
image=np.transpose(color.rgb2lab(image),(0,3,1,2))
image=torch.from_numpy((image+np.array([-50,0,0])[None,:,None,None])).float()
X=torch.unsqueeze(image[:,0,:,:],1).to(device) #set X to the Lightness of the image
image=image[:,1:,:,:].to(device) #image is a and b channel
else:
#convert to grayscale image
#using the matlab formula: 0.2989 * R + 0.5870 * G + 0.1140 * B and load data to gpu
X=(image.clone()*gray).sum(1).to(device).view(-1,1,*in_shape[1:])
image=image.float().to(device)
#----------------------------------------------------------------------------------------
################################### Unet optimization ###################################
#----------------------------------------------------------------------------------------
#clear gradients
optimizer_g.zero_grad()
#generate colorized version with unet
unet_col=None
#print(X.shape,image.shape,classes)
if mode==0:
unet_col=UNet(torch.stack((X,X,X),1)[:,:,0,:,:])
else:
unet_col=UNet(X)
#calculate loss as a function of how good the unet can fool the critic
fooling_loss=criterion(crit(unet_col)[:,0], ones[:batch_size])
#calculate how close the generated pictures are to the ground truth
image_loss=l1loss(unet_col,image)
#combine both losses and weight them
loss_g=fooling_loss+image_loss_weight*image_loss
#backpropagation
loss_g.backward()
optimizer_g.step()
#----------------------------------------------------------------------------------------
################################## Critic optimization ##################################
#----------------------------------------------------------------------------------------
optimizer_c.zero_grad()
real_loss=criterion(crit(image)[:,0],ones[:batch_size])
#requires no gradient in unet col
fake_loss=criterion(crit(unet_col.detach())[:,0],zeros[:batch_size])
loss_c=.5*(real_loss+fake_loss)
loss_c.backward()
optimizer_c.step()
g_running+=loss_g.item()
c_running+=loss_c.item()
loss_hist.append([e,i,loss_g.item(),loss_c.item()])
#report running loss
if (i+len(trainloader)*e)%report_freq==report_freq-1:
print('Epoch %i, batch %i: \tunet loss=%.2e, \tcritic loss=%.2e'%(e+1,i+1,g_running/report_freq,c_running/report_freq))
g_running=0
c_running=0
if s.save_weights and (i+len(trainloader)*e)%save_freq==save_freq-1:
#save parameters
try:
torch.save(UNet.state_dict(),weight_path_ending)
torch.save(crit.state_dict(),crit_path)
except FileNotFoundError:
os.makedirs(weight_path)
torch.save(UNet.state_dict(),weight_path_ending)
torch.save(crit.state_dict(),crit_path)
print("Parameters saved")
if s.save_loss:
#save loss history to file
try:
f=open(loss_path_ending,'a')
np.savetxt(f,loss_hist,'%e')
f.close()
except FileNotFoundError:
os.makedirs(s.loss_path)
np.savetxt(loss_path_ending,loss_hist,'%e')
loss_hist=[]
#update epoch count in dict after each epoch
model_dict[weights_name]["epochs"] = e
#save it to file
try:
with open(model_description_path_ending, "w") as file:
json.dump(model_dict, file, sort_keys=True, indent=4)
except:
print('Could not save to model dictionary (JSON-file)')
def main(argv):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
gray = torch.tensor([0.2989 ,0.5870, 0.1140])[:,None,None].float()
alex = alexnet().to(device)
alex.load_state_dict(torch.load('weights/alexnet.pth'))
alex.eval()
classifier_path = 'weights/alexNorm.pt'
weight_path = None
mbsize = 16
col_space = None
class_net = False
mode = 2
drop_rate = 0
T = .4
classes = 3
stl_path = './stl-10'
help='metric.py -b <batch size> -w <model weight file> -p <2layer classifer weight file> -s <colorspace "rgb" or "lab"> \
-c <use if model is classification model> -m <mode> -d <dropout rate> -t <define temperature> --stl <path to stl dataset>
try:
opts, args = getopt.getopt(argv,"b:w:p:cm:d:t:s:",
["mbsize=",'weight-path=','classifier-path=','c_space=','mode=','drop_rate=','temperature=','stl='])
except getopt.GetoptError:
print(help)
sys.exit(2)
print("opts" ,opts)
for opt, arg in opts:
if opt in ("-b", "--mbsize"):
mbsize = int(arg)
elif opt in ("-w", "--weight-path"):
weight_path = str(arg)
elif opt in ("-p", "--classifier-path"):
classifier_path = str(arg)
elif opt in ("-s","--c_space"):
col_space = str(arg)
elif opt =='-c':
class_net = True
elif opt == '-m':
if arg in ('u','0','unet'):
mode = 0
elif arg in ('color','1','cu'):
mode = 1
elif arg in ('mu','2','middle'):
mode = 2
elif opt in ("-d", "--drop_rate"):
drop_rate = float(arg)
elif opt in ('-t', '--temperature'):
T = float(arg)
elif opt in ('--stl'):
stl_path=arg
if col_space == None:
print('Specify color space')
sys.exit(2)
if class_net:
if col_space == 'yuv':
classes = 42
elif col_space == 'lab':
classes = 150
if mode == 0:
Col_Net = unet(True, drop_rate, classes).to(device)
elif mode == 1:
Col_Net = color_unet(True, drop_rate, classes).to(device)
elif mode == 2:
Col_Net = middle_unet(True, drop_rate, classes).to(device)
Col_Net.load_state_dict(torch.load(weight_path, map_location=device))
Col_Net.eval()
classifier=nn.Sequential(nn.Linear(1000,512),nn.ReLU(),nn.Linear(512,10)).to(device)
classifier.load_state_dict(torch.load(classifier_path, map_location=device))
classifier.eval()
#we trained on the test set and evaluate on the train set since the test set has more labeled images than the trainset
testset = datasets.STL10(stl_path,split='train',transform=transforms.Compose([transforms.ToTensor()]))
testloader = torch.utils.data.DataLoader(testset, batch_size=mbsize, shuffle=True, num_workers=0)
print(pseudo_metric(testloader, col_space, Col_Net, classifier, alex, T))