DissectionProject understand how to drive a GanTester within a dissection project directory structure: it caches data in files, creates image files, and translates data between plain python data types and the pytorch-specific tensors required by GanTester.
| 10 | from io import BytesIO |
| 11 | |
| 12 | class DissectionProject: |
| 13 | ''' |
| 14 | DissectionProject understand how to drive a GanTester within a |
| 15 | dissection project directory structure: it caches data in files, |
| 16 | creates image files, and translates data between plain python data |
| 17 | types and the pytorch-specific tensors required by GanTester. |
| 18 | ''' |
| 19 | def __init__(self, config, project_dir, path_url, public_host): |
| 20 | print('config done', project_dir) |
| 21 | self.use_cuda = torch.cuda.is_available() |
| 22 | self.dissect = config |
| 23 | self.project_dir = project_dir |
| 24 | self.path_url = path_url |
| 25 | self.public_host = public_host |
| 26 | self.cachedir = os.path.join(self.project_dir, 'cache') |
| 27 | self.tester = GanTester( |
| 28 | config.settings, dissectdir=project_dir, |
| 29 | device=torch.device('cuda') if self.use_cuda |
| 30 | else torch.device('cpu')) |
| 31 | self.stdz = [] |
| 32 | |
| 33 | def get_zs(self, size): |
| 34 | if size <= len(self.stdz): |
| 35 | return self.stdz[:size].tolist() |
| 36 | z_tensor = self.tester.standard_z_sample(size) |
| 37 | numpy_z = z_tensor.cpu().numpy() |
| 38 | self.stdz = numpy_z |
| 39 | return self.stdz.tolist() |
| 40 | |
| 41 | def get_z(self, id): |
| 42 | if id < len(self.stdz): |
| 43 | return self.stdz[id] |
| 44 | return self.get_zs((id + 1) * 2)[id] |
| 45 | |
| 46 | def get_zs_for_ids(self, ids): |
| 47 | max_id = max(ids) |
| 48 | if max_id >= len(self.stdz): |
| 49 | self.get_z(max_id) |
| 50 | return self.stdz[ids] |
| 51 | |
| 52 | def get_layers(self): |
| 53 | result = [] |
| 54 | layer_shapes = self.tester.layer_shapes() |
| 55 | for layer in self.tester.layers: |
| 56 | shape = layer_shapes[layer] |
| 57 | result.append(dict( |
| 58 | layer=layer, |
| 59 | channels=shape[1], |
| 60 | shape=[shape[2], shape[3]])) |
| 61 | return result |
| 62 | |
| 63 | def get_units(self, layer): |
| 64 | try: |
| 65 | dlayer = [dl for dl in self.dissect['layers'] |
| 66 | if dl['layer'] == layer][0] |
| 67 | except: |
| 68 | return None |
| 69 |