MCPcopy Create free account
hub / github.com/cjrd/self-supervised-pretraining / main

Function main

utils/plot_basetrain_robust_modified.py:96–217  ·  view source on GitHub ↗
(args)

Source from the content-addressed store, hash-verified

94 return best
95
96def main(args):
97
98 #must pass in dataset arg for proper results
99
100 os.makedirs(args.out_dir, exist_ok=True)
101 # setup plots
102 sns.set_style('darkgrid')
103 sns.set()
104
105 frames = [] #array that collects dataframes from each file
106 if(args.dataset == "all"):
107 dataset_type = "*"
108 else:
109 dataset_type = args.dataset
110
111 nobt_baseline = 0
112
113 #gets all the basetrain results from the specified dataset
114 result_files = glob.glob(os.path.join(args.results_dir, dataset_type + "*basetrain*.json"), recursive=True)
115 result_files2 = glob.glob(os.path.join(args.results_dir, dataset_type + "*bt_robust*.json"), recursive=True)
116 result_files += result_files2
117 result_files.append(os.path.join(args.results_dir, dataset_type + "_results.json")) #gets file with basetrian results
118
119 for resfile in result_files:
120 with open(resfile, 'r') as infile:
121 raw_data = json.load(infile)
122
123 print(resfile)
124
125 types = [] #used for concatenating the resultzs from each basetrained model
126 data = pd.DataFrame(raw_data.values())
127
128
129 #finds the relevant values for moco bt no bt
130 #mainly done to work around the baseline json file (has extra info that we don't want)
131 if (dataset_type + "_results.json") in resfile:
132 linear_data = data[data.result_type=='linear-eval']
133 linear_data = linear_data[linear_data.variant=="linear-eval-lr"]
134
135 data_moco = linear_data[data.basetrain=="moco_v2_800ep"]
136 data_moco = data_moco[data_moco.pretrain_data==dataset_type]
137 data_moco_x = data_moco[(data_moco['pretrain_iters']=="5000")]
138 data_moco_y = data_moco[(data_moco['pretrain_iters']=="0")]
139 data_moco_x = data_moco_x.sort_values(by=['result'],ascending=False).head(1)
140 data_moco_y = data_moco_y.sort_values(by=['result'],ascending=False).head(1)
141 data_moco = pd.concat([data_moco_x,data_moco_y],ignore_index=True)
142
143 #just for asympotote
144 data_nobt = linear_data[data.basetrain=="no"]
145 data_nobt= data_nobt[data_nobt.pretrain_data==dataset_type]
146 nobt_baseline = reduce(data_nobt)
147
148
149 #appended to list
150 types.append(data_moco)
151
152 #all concat to new dataframe
153 data = pd.concat(types, ignore_index=True)

Calls 3

reduceFunction · 0.70
gen_plotsFunction · 0.70
setMethod · 0.45

Tested by

no test coverage detected