def __init__(self, batch_size, latent_dim, head_dim, pretrained):
super(CustomSimCLR18, self).__init__()
self.model = GetModel('Resnet18', pretrained=pretrained)
self.features = torch.nn.Sequential(*list(self.model.children())[:-1])
self.features.add_module('Final Convolution',
nn.Conv2d(512, latent_dim, 1))
self.head = SimCLR_head(latent_dim, head_dim)
self.test_shape(batch_size, latent_dim)
def __init__(self, batch_size, latent_dim, head_dim):
super(CustomSimCLR50, self).__init__()
self.pretrained_model = GetModel('Resnet50')
self.load_weights('./resnet50-1x.pth')
self.features = torch.nn.Sequential(
*list(self.pretrained_model.children())[:-1])
self.features.add_module('Final Convolution',
nn.Conv2d(2048, latent_dim, 1))
self.head = SimCLR_head(latent_dim, head_dim)
self.test_shape(batch_size, latent_dim)
torch.save(checkpoint, opt.out + '/final.pth')
if __name__ == '__main__':
try:
os.makedirs(opt.out)
except IOError:
pass
opt.fid = open(opt.out + '/log.txt', 'w')
print(opt)
print(opt, '\n', file=opt.fid)
dataloader, validloader = GetDataloaders(opt)
net = GetModel(opt)
if opt.log:
opt.writer = SummaryWriter(
comment='_%s_%s' %
(opt.out.replace('\\', '/').split('/')[-1], opt.model))
opt.train_stats = open(
opt.out.replace('\\', '/') + '/train_stats.csv', 'w')
opt.test_stats = open(
opt.out.replace('\\', '/') + '/test_stats.csv', 'w')
print('iter,nsample,time,memory,meanloss', file=opt.train_stats)
print('iter,time,memory,psnr,ssim', file=opt.test_stats)
import time
t0 = time.perf_counter()
if not opt.test:
def main(opt):
opt.device = torch.device(
'cuda' if torch.cuda.is_available() and not opt.cpu else 'cpu')
os.makedirs(opt.out, exist_ok=True)
shutil.copy2('options.py', opt.out)
opt.fid = open(opt.out + '/log.txt', 'w')
ostr = 'ARGS: ' + ' '.join(sys.argv[:])
print(opt, '\n')
print(opt, '\n', file=opt.fid)
print('\n%s\n' % ostr)
print('\n%s\n' % ostr, file=opt.fid)
print('getting dataloader', opt.root)
dataloader, validloader = GetDataloaders(opt)
if opt.log:
opt.writer = SummaryWriter(
log_dir=opt.out,
comment='_%s_%s' %
(opt.out.replace('\\', '/').split('/')[-1], opt.model))
opt.train_stats = open(
opt.out.replace('\\', '/') + '/train_stats.csv', 'w')
opt.test_stats = open(
opt.out.replace('\\', '/') + '/test_stats.csv', 'w')
print('iter,nsample,time,memory,meanloss', file=opt.train_stats)
print('iter,time,memory,psnr,ssim', file=opt.test_stats)
t0 = time.perf_counter()
net = GetModel(opt)
if not opt.test:
train(opt, dataloader, validloader, net)
# torch.save(net.state_dict(), opt.out + '/final.pth')
else:
if len(opt.weights) > 0: # load previous weights?
checkpoint = torch.load(opt.weights)
print('loading checkpoint', opt.weights)
net.load_state_dict(checkpoint['state_dict'])
print('time: %0.1f' % (time.perf_counter() - t0))
testAndMakeCombinedPlots(net, validloader, opt)
opt.fid.close()
if not opt.test:
generate_convergence_plots(opt, opt.out + '/log.txt')
print('time: %0.1f' % (time.perf_counter() - t0))
# optional clean up
if opt.disposableTrainingData and not opt.test:
print('deleting training data')
# preserve a few samples
os.makedirs('%s/training_data_subset' % opt.out, exist_ok=True)
samplecount = 0
for file in glob.glob('%s/*' % opt.root):
if os.path.isfile(file):
basename = os.path.basename(file)
shutil.copy2(
file, '%s/training_data_subset/%s' % (opt.out, basename))
samplecount += 1
if samplecount == 10:
break
shutil.rmtree(opt.root)